Repository: microsoft/MLOS

Issue:
mlos_bench: read results file format tests
@poojanilangekar, could you please add some tests for the changes in #460 so that we make sure that the code accepts both input file formats we expect to be able to support?  Thanks!


Discussion:

Resolving patch:
diff --git a/mlos_bench/mlos_bench/environments/local/local_env.py b/mlos_bench/mlos_bench/environments/local/local_env.py
--- a/mlos_bench/mlos_bench/environments/local/local_env.py
+++ b/mlos_bench/mlos_bench/environments/local/local_env.py
@@ -179,22 +179,36 @@ def run(self) -> Tuple[Status, Optional[Dict[str, float]]]:
             _LOG.debug("Not reading the data at: %s", self)
             return (Status.SUCCEEDED, {})
 
-        data: pandas.DataFrame = pandas.read_csv(
+        data = self._normalize_columns(pandas.read_csv(
             self._config_loader_service.resolve_path(
-                self._read_results_file, extra_paths=[self._temp_dir]))
-
-        if sys.platform == 'win32':
-            data.rename(str.rstrip, axis='columns', inplace=True)
+                self._read_results_file, extra_paths=[self._temp_dir]),
+            index_col=False,
+        ))
 
         _LOG.debug("Read data:\n%s", data)
         if list(data.columns) == ["metric", "value"]:
-            _LOG.warning(
-                "Local run has %d rows: assume long format of (metric, value)", len(data))
+            _LOG.info("Local results have (metric,value) header and %d rows: assume long format", len(data))
             data = pandas.DataFrame([data.value.to_list()], columns=data.metric.to_list())
+        elif len(data) == 1:
+            _LOG.info("Local results have 1 row: assume wide format")
+        else:
+            raise ValueError(f"Invalid data format: {data}")
 
         data_dict = data.iloc[-1].to_dict()
         _LOG.info("Local run complete: %s ::\n%s", self, data_dict)
-        return (Status.SUCCEEDED, data_dict) if data_dict else (Status.FAILED, None)
+        return (Status.SUCCEEDED, data_dict)
+
+    @staticmethod
+    def _normalize_columns(data: pandas.DataFrame) -> pandas.DataFrame:
+        """
+        Strip trailing spaces from column names (Windows only).
+        """
+        # Windows cmd interpretation of > redirect symbols can leave trailing spaces in
+        # the final column, which leads to misnamed columns.
+        # For now, we simply strip trailing spaces from column names to account for that.
+        if sys.platform == 'win32':
+            data.rename(str.rstrip, axis='columns', inplace=True)
+        return data
 
     def status(self) -> Tuple[Status, List[Tuple[datetime, str, Any]]]:
 
@@ -204,25 +218,25 @@ def status(self) -> Tuple[Status, List[Tuple[datetime, str, Any]]]:
 
         assert self._temp_dir is not None
         try:
+            fname = self._config_loader_service.resolve_path(
+                self._read_telemetry_file, extra_paths=[self._temp_dir])
+
             # FIXME: We should not be assuming that the only output file type is a CSV.
-            data: pandas.DataFrame = pandas.read_csv(
-                self._config_loader_service.resolve_path(
-                    self._read_telemetry_file, extra_paths=[self._temp_dir]),
-                parse_dates=[0],
-            )
+            data = self._normalize_columns(
+                pandas.read_csv(fname, index_col=False, parse_dates=[0]))
+
+            expected_col_names = ["timestamp", "metric", "value"]
+            if len(data.columns) != len(expected_col_names):
+                raise ValueError(f'Telemetry data must have columns {expected_col_names}')
+            elif list(data.columns) != expected_col_names:
+                # Assume no header - this is ok for telemetry data.
+                data = pandas.read_csv(
+                    fname, index_col=False, parse_dates=[0], names=expected_col_names)
         except FileNotFoundError as ex:
             _LOG.warning("Telemetry CSV file not found: %s :: %s", self._read_telemetry_file, ex)
             return (status, [])
 
-        if sys.platform == 'win32':
-            data.rename(str.rstrip, axis='columns', inplace=True)
-
         _LOG.debug("Read telemetry data:\n%s", data)
-        if list(data.columns) != ["timestamp", "metric", "value"]:
-            _LOG.warning(
-                'Telemetry CSV file should have columns ["timestamp", "metric", "value"] :: %s',
-                self._read_telemetry_file)
-
         col_dtypes: Mapping[int, Type] = {0: datetime}
         return (status, [
             (pandas.Timestamp(ts).to_pydatetime(), metric, value)


Test patch:
diff --git a/mlos_bench/mlos_bench/tests/environments/local_env_test.py b/mlos_bench/mlos_bench/tests/environments/local_env_test.py
--- a/mlos_bench/mlos_bench/tests/environments/local_env_test.py
+++ b/mlos_bench/mlos_bench/tests/environments/local_env_test.py
@@ -3,10 +3,12 @@
 # Licensed under the MIT License.
 #
 """
-Unit tests for mock benchmark environment.
+Unit tests for LocalEnv benchmark environment.
 """
 from datetime import datetime, timedelta
 
+import pytest
+
 from mlos_bench.environments.local.local_env import LocalEnv
 from mlos_bench.services.config_persistence import ConfigPersistenceService
 from mlos_bench.services.local.local_exec import LocalExecService
@@ -66,3 +68,196 @@ def test_local_env(tunable_groups: TunableGroups) -> None:
             (ts2, "cpu_load", 0.8),
             (ts2, "mem_usage", 20480.0),
         ]
+
+
+def test_local_env_telemetry_no_header(tunable_groups: TunableGroups) -> None:
+    """
+    Read the telemetry data with no header.
+    """
+    ts1 = datetime.utcnow()
+    ts1 -= timedelta(microseconds=ts1.microsecond)  # Round to a second
+    ts2 = ts1 + timedelta(minutes=1)
+
+    time_str1 = ts1.strftime("%Y-%m-%d %H:%M:%S")
+    time_str2 = ts2.strftime("%Y-%m-%d %H:%M:%S")
+
+    local_env = LocalEnv(
+        name="Test Local Env",
+        config={
+            "run": [
+                f"echo {time_str1},cpu_load,0.65 > telemetry.csv",
+                f"echo {time_str1},mem_usage,10240 >> telemetry.csv",
+                f"echo {time_str2},cpu_load,0.8 >> telemetry.csv",
+                f"echo {time_str2},mem_usage,20480 >> telemetry.csv",
+            ],
+            "read_telemetry_file": "telemetry.csv",
+        },
+        service=LocalExecService(parent=ConfigPersistenceService()),
+    )
+
+    with local_env as env_context:
+
+        assert env_context.setup(tunable_groups)
+        (status, _data) = env_context.run()
+        assert status.is_succeeded()
+
+        (status, telemetry) = env_context.status()
+        assert status.is_good()
+        assert telemetry == [
+            (ts1, "cpu_load", 0.65),
+            (ts1, "mem_usage", 10240.0),
+            (ts2, "cpu_load", 0.8),
+            (ts2, "mem_usage", 20480.0),
+        ]
+
+
+def test_local_env_telemetry_wrong_header(tunable_groups: TunableGroups) -> None:
+    """
+    Read the telemetry data with incorrect header.
+    """
+    ts1 = datetime.utcnow()
+    ts1 -= timedelta(microseconds=ts1.microsecond)  # Round to a second
+    ts2 = ts1 + timedelta(minutes=1)
+
+    time_str1 = ts1.strftime("%Y-%m-%d %H:%M:%S")
+    time_str2 = ts2.strftime("%Y-%m-%d %H:%M:%S")
+
+    local_env = LocalEnv(
+        name="Test Local Env",
+        config={
+            "run": [
+                # Error: the data is correct, but the header has unexpected column names
+                "echo 'ts,metric_name,metric_value' > telemetry.csv",
+                f"echo {time_str1},cpu_load,0.65 >> telemetry.csv",
+                f"echo {time_str1},mem_usage,10240 >> telemetry.csv",
+                f"echo {time_str2},cpu_load,0.8 >> telemetry.csv",
+                f"echo {time_str2},mem_usage,20480 >> telemetry.csv",
+            ],
+            "read_telemetry_file": "telemetry.csv",
+        },
+        service=LocalExecService(parent=ConfigPersistenceService()),
+    )
+
+    with local_env as env_context:
+
+        assert env_context.setup(tunable_groups)
+        (status, _data) = env_context.run()
+        assert status.is_succeeded()
+
+        with pytest.raises(ValueError):
+            env_context.status()
+
+
+def test_local_env_telemetry_invalid(tunable_groups: TunableGroups) -> None:
+    """
+    Fail when the telemetry data has wrong format.
+    """
+    ts1 = datetime.utcnow()
+    ts1 -= timedelta(microseconds=ts1.microsecond)  # Round to a second
+    ts2 = ts1 + timedelta(minutes=1)
+
+    time_str1 = ts1.strftime("%Y-%m-%d %H:%M:%S")
+    time_str2 = ts2.strftime("%Y-%m-%d %H:%M:%S")
+
+    local_env = LocalEnv(
+        name="Test Local Env",
+        config={
+            "run": [
+                # Error: too many columns
+                f"echo {time_str1},EXTRA,cpu_load,0.65 > telemetry.csv",
+                f"echo {time_str1},EXTRA,mem_usage,10240 >> telemetry.csv",
+                f"echo {time_str2},EXTRA,cpu_load,0.8 >> telemetry.csv",
+                f"echo {time_str2},EXTRA,mem_usage,20480 >> telemetry.csv",
+            ],
+            "read_telemetry_file": "telemetry.csv",
+        },
+        service=LocalExecService(parent=ConfigPersistenceService()),
+    )
+
+    with local_env as env_context:
+
+        assert env_context.setup(tunable_groups)
+        (status, _data) = env_context.run()
+        assert status.is_succeeded()
+
+        with pytest.raises(ValueError):
+            env_context.status()
+
+
+def test_local_env_telemetry_invalid_ts(tunable_groups: TunableGroups) -> None:
+    """
+    Fail when the telemetry data has wrong format.
+    """
+    local_env = LocalEnv(
+        name="Test Local Env",
+        config={
+            "run": [
+                # Error: field 1 must be a timestamp
+                "echo 1,cpu_load,0.65 > telemetry.csv",
+                "echo 2,mem_usage,10240 >> telemetry.csv",
+                "echo 3,cpu_load,0.8 >> telemetry.csv",
+                "echo 4,mem_usage,20480 >> telemetry.csv",
+            ],
+            "read_telemetry_file": "telemetry.csv",
+        },
+        service=LocalExecService(parent=ConfigPersistenceService()),
+    )
+
+    with local_env as env_context:
+
+        assert env_context.setup(tunable_groups)
+        (status, _data) = env_context.run()
+        assert status.is_succeeded()
+
+        with pytest.raises(ValueError):
+            env_context.status()
+
+
+def test_local_env_results_no_header(tunable_groups: TunableGroups) -> None:
+    """
+    Fail if the results are not in the expected format.
+    """
+    local_env = LocalEnv(
+        name="Test Local Env",
+        config={
+            "run": [
+                # No header
+                "echo 'latency,10' > output.csv",
+                "echo 'throughput,66' >> output.csv",
+                "echo 'score,0.9' >> output.csv",
+            ],
+            "read_results_file": "output.csv",
+        },
+        service=LocalExecService(parent=ConfigPersistenceService()),
+    )
+    with local_env as env_context:
+        assert env_context.setup(tunable_groups)
+        with pytest.raises(ValueError):
+            env_context.run()
+
+
+def test_local_env_wide(tunable_groups: TunableGroups) -> None:
+    """
+    Produce benchmark data in wide format and read it.
+    """
+    local_env = LocalEnv(
+        name="Test Local Env",
+        config={
+            "run": [
+                "echo 'latency,throughput,score' > output.csv",
+                "echo '10,66,0.9' >> output.csv",
+            ],
+            "read_results_file": "output.csv",
+        },
+        tunables=tunable_groups,
+        service=LocalExecService(parent=ConfigPersistenceService()),
+    )
+    with local_env as env_context:
+        assert env_context.setup(tunable_groups)
+        (status, data) = env_context.run()
+        assert status.is_succeeded()
+        assert data == {
+            "latency": 10.0,
+            "throughput": 66.0,
+            "score": 0.9,
+        }