start using datetimes with fractional seconds with mysql

Author: bpkroth

mlos_bench/mlos_bench/storage/sql/experiment.py

@@ -367,11 +367,7 @@ def _new_trial(
         ts_start: datetime | None = None,
         config: dict[str, Any] | None = None,
     ) -> Storage.Trial:
-        # MySQL can round microseconds into the future causing scheduler to skip trials.
-        # Truncate microseconds to avoid this issue.
-        ts_start = utcify_timestamp(ts_start or datetime.now(UTC), origin="local").replace(
-            microsecond=0
-        )
+        ts_start = utcify_timestamp(ts_start or datetime.now(UTC), origin="local")
         _LOG.debug("Create trial: %s:%d @ %s", self._experiment_id, self._trial_id, ts_start)
         with self._engine.begin() as conn:
             try:

mlos_bench/mlos_bench/storage/sql/schema.py

@@ -39,6 +39,7 @@
     create_mock_engine,
     inspect,
 )
+from sqlalchemy.dialects import mysql
 from sqlalchemy.engine import Engine
 
 from mlos_bench.util import path_join
@@ -104,8 +105,8 @@ def __init__(self, engine: Engine | None):
             Column("git_repo", String(1024), nullable=False),
             Column("git_commit", String(40), nullable=False),
             # For backwards compatibility, we allow NULL for ts_start.
-            Column("ts_start", DateTime),
-            Column("ts_end", DateTime),
+            Column("ts_start", DateTime().with_variant(mysql.DATETIME(fsp=6), "mysql")),
+            Column("ts_end", DateTime().with_variant(mysql.DATETIME(fsp=6), "mysql")),
             # Should match the text IDs of `mlos_bench.environments.Status` enum:
             # For backwards compatibility, we allow NULL for status.
             Column("status", String(self._STATUS_LEN)),
@@ -179,8 +180,16 @@ def __init__(self, engine: Engine | None):
             Column("trial_id", Integer, nullable=False),
             Column("config_id", Integer, nullable=False),
             Column("trial_runner_id", Integer, nullable=True, default=None),
-            Column("ts_start", DateTime, nullable=False),
-            Column("ts_end", DateTime),
+            Column(
+                "ts_start",
+                DateTime().with_variant(mysql.DATETIME(fsp=6), "mysql"),
+                nullable=False,
+            ),
+            Column(
+                "ts_end",
+                DateTime().with_variant(mysql.DATETIME(fsp=6), "mysql"),
+                nullable=True,
+            ),
             # Should match the text IDs of `mlos_bench.environments.Status` enum:
             Column("status", String(self._STATUS_LEN), nullable=False),
             PrimaryKeyConstraint("exp_id", "trial_id"),
@@ -232,7 +241,12 @@ def __init__(self, engine: Engine | None):
             self._meta,
             Column("exp_id", String(self._ID_LEN), nullable=False),
             Column("trial_id", Integer, nullable=False),
-            Column("ts", DateTime(timezone=True), nullable=False, default="now"),
+            Column(
+                "ts",
+                DateTime(timezone=True).with_variant(mysql.DATETIME(fsp=6), "mysql"),
+                nullable=False,
+                default="now",
+            ),
             Column("status", String(self._STATUS_LEN), nullable=False),
             UniqueConstraint("exp_id", "trial_id", "ts"),
             ForeignKeyConstraint(
@@ -267,7 +281,12 @@ def __init__(self, engine: Engine | None):
             self._meta,
             Column("exp_id", String(self._ID_LEN), nullable=False),
             Column("trial_id", Integer, nullable=False),
-            Column("ts", DateTime(timezone=True), nullable=False, default="now"),
+            Column(
+                "ts",
+                DateTime(timezone=True).with_variant(mysql.DATETIME(fsp=6), "mysql"),
+                nullable=False,
+                default="now",
+            ),
             Column("metric_id", String(self._ID_LEN), nullable=False),
             Column("metric_value", String(self._METRIC_VALUE_LEN)),
             UniqueConstraint("exp_id", "trial_id", "ts", "metric_id"),