ray-project--ray/python/ray/train/v2/tests/test_state_export.py

import json
import os

import pytest

import ray
from ray.train.v2._internal.state.schema import (
    RunAttemptStatus,
    RunStatus,
)
from ray.train.v2._internal.state.state_actor import get_or_create_state_actor
from ray.train.v2.tests.util import (
    create_mock_train_run,
    create_mock_train_run_attempt,
)


@pytest.fixture
def shutdown_only():
    yield
    ray.shutdown()


def _get_export_file_path() -> str:
    return os.path.join(
        ray._private.worker._global_node.get_session_dir_path(),
        "logs",
        "export_events",
        "event_EXPORT_TRAIN_STATE.log",
    )


def _get_exported_data():
    exported_file = _get_export_file_path()
    assert os.path.isfile(exported_file)

    with open(exported_file, "r") as f:
        data = f.readlines()

    return [json.loads(line) for line in data]


@pytest.fixture
def enable_export_api_config(shutdown_only):
    """Enable export API for the EXPORT_TRAIN_RUN resource type."""
    ray.init(
        num_cpus=4,
        runtime_env={
            "env_vars": {"RAY_enable_export_api_write_config": "EXPORT_TRAIN_RUN"}
        },
    )
    yield


@pytest.fixture
def enable_export_api_write(shutdown_only):
    """Enable export API for all resource types."""
    ray.init(
        num_cpus=4,
        runtime_env={"env_vars": {"RAY_enable_export_api_write": "1"}},
    )
    yield


def test_export_disabled(ray_start_4_cpus):
    """Test that no export files are created when export API is disabled."""
    state_actor = get_or_create_state_actor()

    # Create or update train run
    ray.get(state_actor.create_or_update_train_run.remote(create_mock_train_run()))
    ray.get(
        state_actor.create_or_update_train_run_attempt.remote(
            create_mock_train_run_attempt()
        )
    )

    # Check that no export files were created
    assert not os.path.exists(_get_export_file_path())


def _test_train_run_export():
    """Test that train run export events are written when export API is enabled."""
    state_actor = get_or_create_state_actor()

    # Create or update train run
    ray.get(
        state_actor.create_or_update_train_run.remote(
            create_mock_train_run(RunStatus.RUNNING)
        )
    )

    # Check that export files were created
    data = _get_exported_data()

    assert len(data) == 1
    assert data[0]["source_type"] == "EXPORT_TRAIN_RUN"
    assert data[0]["event_data"]["status"] == "RUNNING"


def test_export_train_run_enabled_by_config(enable_export_api_config):
    _test_train_run_export()


def test_export_train_run(enable_export_api_write):
    _test_train_run_export()


def test_export_train_run_run_settings_fields(enable_export_api_write):
    """Test that RunSettings fields are exported with expected shapes/types.

    This is a regression test for the JSON export format produced by the Train state
    export logger and the `export_train_state.proto` schema.
    """
    state_actor = get_or_create_state_actor()

    ray.get(
        state_actor.create_or_update_train_run.remote(
            create_mock_train_run(RunStatus.RUNNING)
        )
    )

    data = _get_exported_data()
    assert len(data) == 1
    run_settings = data[0]["event_data"]["run_settings"]

    assert run_settings == {
        "backend_config": {
            "config": {},
            "framework": "TRAINING_FRAMEWORK_UNSPECIFIED",
        },
        "scaling_config": {
            "num_workers_fixed": 1,
            "placement_strategy": "PACK",
            "use_gpu": False,
            "use_tpu": False,
        },
        "datasets": ["dataset_1"],
        "data_config": {
            "all": {},
            "enable_shard_locality": True,
            "data_execution_options": {
                "default": {
                    "actor_locality_enabled": True,
                    "exclude_resources": {
                        "CPU": 0.0,
                        "GPU": 0.0,
                        "memory": 0.0,
                        "object_store_memory": 0.0,
                    },
                    "preserve_order": False,
                    "resource_limits": {
                        "CPU": "inf",
                        "GPU": "inf",
                        "memory": "inf",
                        "object_store_memory": "inf",
                    },
                    "verbose_progress": True,
                },
                "per_dataset_execution_options": {},
            },
        },
        "run_config": {
            "name": "test_run",
            "failure_config": {"controller_failure_limit": -1, "max_failures": 0},
            "worker_runtime_env": {"type": "conda"},
            "checkpoint_config": {"checkpoint_score_order": "MAX"},
            "storage_path": "s3://bucket/path",
        },
    }


def test_export_oneof_num_workers(enable_export_api_write):
    """
    Test that the num_workers oneof field exports correctly for both fixed for a fixed
    number of workers and range for when elastic training is enabled.
    """
    state_actor = get_or_create_state_actor()

    run_fixed = create_mock_train_run(RunStatus.RUNNING, id="fixed_workers")
    run_fixed.run_settings.scaling_config.num_workers = 4

    run_range = create_mock_train_run(RunStatus.RUNNING, id="range_workers")
    run_range.run_settings.scaling_config.num_workers = (2, 8)

    ray.get(
        [
            state_actor.create_or_update_train_run.remote(run_fixed),
            state_actor.create_or_update_train_run.remote(run_range),
        ]
    )

    runs = _get_exported_data()
    assert len(runs) == 2

    runs_by_id = {run["event_data"]["id"]: run for run in runs}

    # Fixed num_workers
    fixed_scaling = runs_by_id["fixed_workers"]["event_data"]["run_settings"][
        "scaling_config"
    ]
    assert fixed_scaling["num_workers_fixed"] == 4
    assert "num_workers_range" not in fixed_scaling

    # Range num_workers
    range_scaling = runs_by_id["range_workers"]["event_data"]["run_settings"][
        "scaling_config"
    ]
    assert range_scaling["num_workers_range"] == {"min": 2, "max": 8}
    assert "num_workers_fixed" not in range_scaling


def test_export_train_loop_config_integer_keys(enable_export_api_write):
    """Test that integer keys in train_loop_config are properly stringified in export."""
    state_actor = get_or_create_state_actor()

    ray.get(
        state_actor.create_or_update_train_run.remote(
            create_mock_train_run(train_loop_config={1: "one", 2: "two"})
        )
    )

    data = _get_exported_data()
    assert len(data) == 1

    run_settings = data[0]["event_data"]["run_settings"]

    assert run_settings["train_loop_config"] == {"1": "one", "2": "two"}


def test_export_train_run_attempt(enable_export_api_write):
    """Test that train run attempt export events are written when export API is enabled."""
    state_actor = get_or_create_state_actor()

    # Create or update train run attempt
    ray.get(
        state_actor.create_or_update_train_run_attempt.remote(
            create_mock_train_run_attempt(RunAttemptStatus.RUNNING)
        )
    )

    data = _get_exported_data()
    assert len(data) == 1
    assert data[0]["source_type"] == "EXPORT_TRAIN_RUN_ATTEMPT"
    assert data[0]["event_data"]["status"] == "RUNNING"


def test_export_oneof_datasets_to_split(enable_export_api_write):
    """Test that proto oneof fields are exported with only the active variant set.

    In ExportTrainRunEventData.RunSettings.DataConfig, `datasets_to_split` is a oneof:
    - `all`
    - `datasets` (StringList)
    """
    state_actor = get_or_create_state_actor()

    run_all = create_mock_train_run(RunStatus.RUNNING, id="with_all")
    run_all.run_settings.data_config.datasets_to_split = "all"

    run_ds = create_mock_train_run(RunStatus.RUNNING, id="with_datasets")
    run_ds.run_settings.data_config.datasets_to_split = ["dataset_1", "dataset_2"]

    ray.get(
        [
            state_actor.create_or_update_train_run.remote(run_all),
            state_actor.create_or_update_train_run.remote(run_ds),
        ]
    )

    runs = _get_exported_data()
    assert len(runs) == 2

    runs_by_id = {run["event_data"]["id"]: run for run in runs}
    assert set(runs_by_id.keys()) == {"with_all", "with_datasets"}

    # Verify train data config when splitting all datasets
    all_settings = runs_by_id["with_all"]["event_data"]["run_settings"]
    assert "data_config" in all_settings
    assert "all" in all_settings["data_config"]
    assert "datasets" not in all_settings["data_config"]

    # Verify train data config when splitting specific datasets
    ds_settings = runs_by_id["with_datasets"]["event_data"]["run_settings"]
    assert "data_config" in ds_settings
    assert "datasets" in ds_settings["data_config"]
    assert (
        ds_settings["data_config"]["datasets"]["values"]
        == run_ds.run_settings.data_config.datasets_to_split
    )
    assert "all" not in ds_settings["data_config"]


def test_export_oneof_bundle_label_selector(enable_export_api_write):
    """Test that proto oneof fields are exported with only the active variant set.

    In ExportTrainRunEventData.RunSettings.ScalingConfig, `bundle_label_selector` is a oneof:
    - `label_selector_single` (StringMap)
    - `label_selector_list` (StringMapList)
    """
    state_actor = get_or_create_state_actor()

    run_single = create_mock_train_run(RunStatus.RUNNING, id="with_single")
    run_single.run_settings.scaling_config.bundle_label_selector = {"k": "v"}

    run_list = create_mock_train_run(RunStatus.RUNNING, id="with_list")
    run_list.run_settings.scaling_config.bundle_label_selector = [
        {"k1": "v1"},
        {"k2": "v2"},
    ]

    run_none = create_mock_train_run(RunStatus.RUNNING, id="with_none")

    ray.get(
        [
            state_actor.create_or_update_train_run.remote(run_single),
            state_actor.create_or_update_train_run.remote(run_list),
            state_actor.create_or_update_train_run.remote(run_none),
        ]
    )

    runs = _get_exported_data()
    assert len(runs) == 3

    runs_by_id = {run["event_data"]["id"]: run for run in runs}
    assert set(runs_by_id.keys()) == {"with_single", "with_list", "with_none"}

    # Verify single dict selector
    single_scaling = runs_by_id["with_single"]["event_data"]["run_settings"][
        "scaling_config"
    ]
    assert single_scaling["label_selector_single"] == {"values": {"k": "v"}}
    assert "label_selector_list" not in single_scaling

    # Verify list of dicts selector
    list_scaling = runs_by_id["with_list"]["event_data"]["run_settings"][
        "scaling_config"
    ]
    assert list_scaling["label_selector_list"] == {
        "values": [{"values": {"k1": "v1"}}, {"values": {"k2": "v2"}}]
    }
    assert "label_selector_single" not in list_scaling

    # Verify unset selector
    none_scaling = runs_by_id["with_none"]["event_data"]["run_settings"][
        "scaling_config"
    ]
    assert "label_selector_single" not in none_scaling
    assert "label_selector_list" not in none_scaling


def test_export_multiple_source_types(enable_export_api_write):
    """Test that multiple source types (Run and RunAttempt) can be written to the same file."""
    state_actor = get_or_create_state_actor()

    events = [
        state_actor.create_or_update_train_run.remote(
            create_mock_train_run(RunStatus.RUNNING)
        ),
        state_actor.create_or_update_train_run_attempt.remote(
            create_mock_train_run_attempt(
                attempt_id="attempt_1", status=RunAttemptStatus.RUNNING
            )
        ),
        state_actor.create_or_update_train_run_attempt.remote(
            create_mock_train_run_attempt(
                attempt_id="attempt_2", status=RunAttemptStatus.RUNNING
            )
        ),
        state_actor.create_or_update_train_run_attempt.remote(
            create_mock_train_run_attempt(
                attempt_id="attempt_1", status=RunAttemptStatus.FINISHED
            )
        ),
        state_actor.create_or_update_train_run_attempt.remote(
            create_mock_train_run_attempt(
                attempt_id="attempt_2", status=RunAttemptStatus.FINISHED
            )
        ),
        state_actor.create_or_update_train_run.remote(
            create_mock_train_run(RunStatus.FINISHED)
        ),
    ]
    ray.get(events)

    data = _get_exported_data()
    assert len(data) == len(events)

    expected_source_types = (
        ["EXPORT_TRAIN_RUN"] + ["EXPORT_TRAIN_RUN_ATTEMPT"] * 4 + ["EXPORT_TRAIN_RUN"]
    )
    expected_statuses = ["RUNNING"] * 3 + ["FINISHED"] * 3

    assert [d["source_type"] for d in data] == expected_source_types
    assert [d["event_data"]["status"] for d in data] == expected_statuses


def test_export_execution_options_with_inf_resource_limits(enable_export_api_write):
    """Test that execution_options containing float('inf') resource limits export cleanly.

    ExecutionResources.for_limits() sets unspecified fields to float('inf'), so any
    DataConfig with per-dataset execution options will contain inf values. These are
    non-standard in JSON and caused MessageToDict to raise ValueError during export,
    crashing the send_event call entirely.
    """
    from ray.data._internal.execution.interfaces.execution_options import (
        ExecutionOptions,
    )
    from ray.train.v2._internal.state.schema import DataExecutionOptions
    from ray.train.v2._internal.state.util import execution_options_to_model

    state_actor = get_or_create_state_actor()

    run = create_mock_train_run(RunStatus.RUNNING)
    # Default ExecutionOptions uses for_limits(), setting all resource fields to inf.
    default_model = execution_options_to_model(ExecutionOptions())
    run.run_settings.data_config.data_execution_options = DataExecutionOptions(
        default=default_model,
        per_dataset_execution_options={"train": default_model},
    )

    # Must not raise — previously MessageToDict crashed on inf values.
    ray.get(state_actor.create_or_update_train_run.remote(run))

    data = _get_exported_data()
    assert len(data) == 1

    exported_exec_opts = data[0]["event_data"]["run_settings"]["data_config"][
        "data_execution_options"
    ]
    # inf must be serialized as the string "inf", not a bare float, in both
    # the default and per-dataset override slots.
    assert exported_exec_opts["default"]["resource_limits"]["CPU"] == "inf"
    assert exported_exec_opts["default"]["resource_limits"]["GPU"] == "inf"
    assert (
        exported_exec_opts["per_dataset_execution_options"]["train"]["resource_limits"][
            "CPU"
        ]
        == "inf"
    )


def test_export_per_dataset_execution_options(enable_export_api_write):
    """Test that a fully populated per_dataset_execution_options round-trips through export.

    Each dataset's ExecutionOptions has distinct values so we can confirm none
    of them collapse onto the default or each other during proto serialization.
    """
    from ray.train.v2._internal.state.schema import (
        DataExecutionOptions,
        ExecutionOptions as ExecutionOptionsSchema,
    )

    state_actor = get_or_create_state_actor()

    default = ExecutionOptionsSchema(
        resource_limits={"CPU": 1.0, "GPU": 0.0},
        exclude_resources={"CPU": 0.0, "GPU": 0.0},
        preserve_order=False,
        actor_locality_enabled=True,
        verbose_progress=True,
    )
    train_opts = ExecutionOptionsSchema(
        resource_limits={"CPU": 8.0, "GPU": 2.0},
        exclude_resources={"CPU": 1.0, "GPU": 0.0},
        preserve_order=True,
        actor_locality_enabled=False,
        verbose_progress=False,
    )
    eval_opts = ExecutionOptionsSchema(
        resource_limits={"CPU": 4.0, "GPU": 1.0},
        exclude_resources={"CPU": 0.5, "GPU": 0.0},
        preserve_order=False,
        actor_locality_enabled=True,
        verbose_progress=False,
    )

    run = create_mock_train_run(RunStatus.RUNNING)
    run.run_settings.data_config.data_execution_options = DataExecutionOptions(
        default=default,
        per_dataset_execution_options={"train": train_opts, "eval": eval_opts},
    )

    ray.get(state_actor.create_or_update_train_run.remote(run))

    data = _get_exported_data()
    assert len(data) == 1

    exported = data[0]["event_data"]["run_settings"]["data_config"][
        "data_execution_options"
    ]
    assert exported == {
        "default": default.dict(),
        "per_dataset_execution_options": {
            "train": train_opts.dict(),
            "eval": eval_opts.dict(),
        },
    }


def test_export_optional_fields(enable_export_api_write):
    """Test that optional fields are correctly exported when present and absent.

    This covers both top-level optional fields on TrainRun/TrainRunAttempt and
    optional nested fields inside TrainRun.run_settings (e.g., scaling/data/run
    configs).
    """
    state_actor = get_or_create_state_actor()

    # Create run with optional fields
    run_with_optional = create_mock_train_run(RunStatus.FINISHED)
    run_with_optional.status_detail = "Finished with details"
    run_with_optional.end_time_ns = 1000000000000000000
    run_with_optional.run_settings.train_loop_config = {"epochs": 1}
    run_with_optional.run_settings.scaling_config.resources_per_worker = {"CPU": 1}
    run_with_optional.run_settings.scaling_config.accelerator_type = "A100"
    run_with_optional.run_settings.scaling_config.topology = "v4-8"
    run_with_optional.run_settings.scaling_config.bundle_label_selector = {"k": "v"}
    run_with_optional.run_settings.data_config.datasets_to_split = ["dataset_1"]
    run_with_optional.run_settings.run_config.checkpoint_config.num_to_keep = 2
    run_with_optional.run_settings.run_config.checkpoint_config.checkpoint_score_attribute = (
        "score"
    )
    run_with_optional.run_settings.run_config.storage_filesystem = "S3FileSystem"

    # Create attempt with optional fields
    attempt_with_optional = create_mock_train_run_attempt(
        attempt_id="attempt_with_optional",
        status=RunAttemptStatus.FINISHED,
    )
    attempt_with_optional.status_detail = "Attempt details"
    attempt_with_optional.end_time_ns = 1000000000000000000

    # Create and update states
    events = [
        state_actor.create_or_update_train_run.remote(create_mock_train_run()),
        state_actor.create_or_update_train_run_attempt.remote(
            create_mock_train_run_attempt()
        ),
        state_actor.create_or_update_train_run.remote(run_with_optional),
        state_actor.create_or_update_train_run_attempt.remote(attempt_with_optional),
    ]
    ray.get(events)

    data = _get_exported_data()

    assert len(data) == 4

    # Verify train run without optional fields
    run_data = data[0]

    assert run_data["source_type"] == "EXPORT_TRAIN_RUN"
    assert "status_detail" not in run_data["event_data"]
    assert "end_time_ns" not in run_data["event_data"]
    assert "run_settings" in run_data["event_data"]
    run_settings = run_data["event_data"]["run_settings"]
    assert "train_loop_config" not in run_settings
    assert "checkpoint_config" in run_settings["run_config"]
    assert "num_to_keep" not in run_settings["run_config"]["checkpoint_config"]
    assert (
        "checkpoint_score_attribute"
        not in run_settings["run_config"]["checkpoint_config"]
    )
    assert "scaling_config" in run_settings
    assert "resources_per_worker" not in run_settings["scaling_config"]
    assert "accelerator_type" not in run_settings["scaling_config"]
    assert "topology" not in run_settings["scaling_config"]
    assert "label_selector_single" not in run_settings["scaling_config"]
    assert "label_selector_list" not in run_settings["scaling_config"]
    assert "data_config" in run_settings
    assert "storage_filesystem" not in run_settings["run_config"]

    # Verify train run attempt without optional fields
    attempt_data = data[1]
    assert attempt_data["source_type"] == "EXPORT_TRAIN_RUN_ATTEMPT"
    assert "status_detail" not in attempt_data["event_data"]
    assert "end_time_ns" not in attempt_data["event_data"]

    # Verify train run with optional fields
    run_data = data[2]
    assert run_data["source_type"] == "EXPORT_TRAIN_RUN"
    assert run_data["event_data"]["status_detail"] == "Finished with details"
    assert "end_time_ns" in run_data["event_data"]

    run_settings = run_data["event_data"]["run_settings"]
    assert run_settings["train_loop_config"] == {"epochs": 1.0}
    assert run_settings["scaling_config"]["resources_per_worker"]["values"] == {
        k: float(v)
        for k, v in run_with_optional.run_settings.scaling_config.resources_per_worker.items()
    }
    assert (
        run_settings["scaling_config"]["accelerator_type"]
        == run_with_optional.run_settings.scaling_config.accelerator_type
    )
    assert (
        run_settings["scaling_config"]["topology"]
        == run_with_optional.run_settings.scaling_config.topology
    )
    assert run_settings["scaling_config"]["label_selector_single"] == {
        "values": {"k": "v"}
    }
    assert (
        run_settings["data_config"]["datasets"]["values"]
        == run_with_optional.run_settings.data_config.datasets_to_split
    )
    assert run_settings["run_config"]["checkpoint_config"]["num_to_keep"] == str(
        run_with_optional.run_settings.run_config.checkpoint_config.num_to_keep
    )
    assert (
        run_settings["run_config"]["checkpoint_config"]["checkpoint_score_attribute"]
        == run_with_optional.run_settings.run_config.checkpoint_config.checkpoint_score_attribute
    )
    assert (
        run_settings["run_config"]["storage_filesystem"]
        == run_with_optional.run_settings.run_config.storage_filesystem
    )

    # Verify train run attempt with optional fields
    attempt_data = data[3]
    assert attempt_data["source_type"] == "EXPORT_TRAIN_RUN_ATTEMPT"
    assert attempt_data["event_data"]["status_detail"] == "Attempt details"
    assert "end_time_ns" in attempt_data["event_data"]


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))