star

venv/lib/python3.9/site-packages/mlflow/store/tracking/dbmodels/initial_models.py

@@ -0,0 +1,243 @@
+# Snapshot of MLflow DB models as of the 0.9.1 release, prior to the first database migration.
+# Used to standardize initial database state.
+# Copied with modifications from
+# https://github.com/mlflow/mlflow/blob/v0.9.1/mlflow/store/dbmodels/models.py, which
+# is the first database schema that users could be running. In particular, modifications have
+# been made to substitute constants from MLflow with hard-coded values (e.g. replacing
+# SourceType.to_string(SourceType.NOTEBOOK) with the constant "NOTEBOOK") and ensure
+# that all constraint names are unique. Note that pre-1.0 database schemas did not have unique
+# constraint names - we provided a one-time migration script for pre-1.0 users so that their
+# database schema matched the schema in this file.
+import time
+
+from sqlalchemy import (
+    BigInteger,
+    CheckConstraint,
+    Column,
+    Float,
+    ForeignKey,
+    Integer,
+    PrimaryKeyConstraint,
+    String,
+)
+from sqlalchemy.orm import backref, declarative_base, relationship
+
+Base = declarative_base()
+
+
+SourceTypes = [
+    "NOTEBOOK",
+    "JOB",
+    "LOCAL",
+    "UNKNOWN",
+    "PROJECT",
+]
+
+RunStatusTypes = [
+    "SCHEDULED",
+    "FAILED",
+    "FINISHED",
+    "RUNNING",
+]
+
+
+class SqlExperiment(Base):
+    """
+    DB model for :py:class:`mlflow.entities.Experiment`. These are recorded in ``experiment`` table.
+    """
+
+    __tablename__ = "experiments"
+
+    experiment_id = Column(Integer, autoincrement=True)
+    """
+    Experiment ID: `Integer`. *Primary Key* for ``experiment`` table.
+    """
+    name = Column(String(256), unique=True, nullable=False)
+    """
+    Experiment name: `String` (limit 256 characters). Defined as *Unique* and *Non null* in
+                     table schema.
+    """
+    artifact_location = Column(String(256), nullable=True)
+    """
+    Default artifact location for this experiment: `String` (limit 256 characters). Defined as
+                                                    *Non null* in table schema.
+    """
+    lifecycle_stage = Column(String(32), default="active")
+    """
+    Lifecycle Stage of experiment: `String` (limit 32 characters).
+                                    Can be either ``active`` (default) or ``deleted``.
+    """
+
+    __table_args__ = (
+        CheckConstraint(
+            lifecycle_stage.in_(["active", "deleted"]), name="experiments_lifecycle_stage"
+        ),
+        PrimaryKeyConstraint("experiment_id", name="experiment_pk"),
+    )
+
+    def __repr__(self):
+        return f"<SqlExperiment ({self.experiment_id}, {self.name})>"
+
+
+class SqlRun(Base):
+    """
+    DB model for :py:class:`mlflow.entities.Run`. These are recorded in ``runs`` table.
+    """
+
+    __tablename__ = "runs"
+
+    run_uuid = Column(String(32), nullable=False)
+    """
+    Run UUID: `String` (limit 32 characters). *Primary Key* for ``runs`` table.
+    """
+    name = Column(String(250))
+    """
+    Run name: `String` (limit 250 characters).
+    """
+    source_type = Column(String(20), default="LOCAL")
+    """
+    Source Type: `String` (limit 20 characters). Can be one of ``NOTEBOOK``, ``JOB``, ``PROJECT``,
+                 ``LOCAL`` (default), or ``UNKNOWN``.
+    """
+    source_name = Column(String(500))
+    """
+    Name of source recording the run: `String` (limit 500 characters).
+    """
+    entry_point_name = Column(String(50))
+    """
+    Entry-point name that launched the run run: `String` (limit 50 characters).
+    """
+    user_id = Column(String(256), nullable=True, default=None)
+    """
+    User ID: `String` (limit 256 characters). Defaults to ``null``.
+    """
+    status = Column(String(20), default="SCHEDULED")
+    """
+    Run Status: `String` (limit 20 characters). Can be one of ``RUNNING``, ``SCHEDULED`` (default),
+                ``FINISHED``, ``FAILED``.
+    """
+    start_time = Column(BigInteger, default=int(time.time()))
+    """
+    Run start time: `BigInteger`. Defaults to current system time.
+    """
+    end_time = Column(BigInteger, nullable=True, default=None)
+    """
+    Run end time: `BigInteger`.
+    """
+    source_version = Column(String(50))
+    """
+    Source version: `String` (limit 50 characters).
+    """
+    lifecycle_stage = Column(String(20), default="active")
+    """
+    Lifecycle Stage of run: `String` (limit 32 characters).
+                            Can be either ``active`` (default) or ``deleted``.
+    """
+    artifact_uri = Column(String(200), default=None)
+    """
+    Default artifact location for this run: `String` (limit 200 characters).
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
+    """
+    Experiment ID to which this run belongs to: *Foreign Key* into ``experiment`` table.
+    """
+    experiment = relationship("SqlExperiment", backref=backref("runs", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`.
+    """
+
+    __table_args__ = (
+        CheckConstraint(source_type.in_(SourceTypes), name="source_type"),
+        CheckConstraint(status.in_(RunStatusTypes), name="status"),
+        CheckConstraint(lifecycle_stage.in_(["active", "deleted"]), name="runs_lifecycle_stage"),
+        PrimaryKeyConstraint("run_uuid", name="run_pk"),
+    )
+
+
+class SqlTag(Base):
+    """
+    DB model for :py:class:`mlflow.entities.RunTag`. These are recorded in ``tags`` table.
+    """
+
+    __tablename__ = "tags"
+
+    key = Column(String(250))
+    """
+    Tag key: `String` (limit 250 characters). *Primary Key* for ``tags`` table.
+    """
+    value = Column(String(250), nullable=True)
+    """
+    Value associated with tag: `String` (limit 250 characters). Could be *null*.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this tag belongs to: *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("tags", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    __table_args__ = (PrimaryKeyConstraint("key", "run_uuid", name="tag_pk"),)
+
+    def __repr__(self):
+        return f"<SqlRunTag({self.key}, {self.value})>"
+
+
+class SqlMetric(Base):
+    __tablename__ = "metrics"
+
+    key = Column(String(250))
+    """
+    Metric key: `String` (limit 250 characters). Part of *Primary Key* for ``metrics`` table.
+    """
+    value = Column(Float, nullable=False)
+    """
+    Metric value: `Float`. Defined as *Non-null* in schema.
+    """
+    timestamp = Column(BigInteger, default=lambda: int(time.time()))
+    """
+    Timestamp recorded for this metric entry: `BigInteger`. Part of *Primary Key* for
+                                               ``metrics`` table.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``metrics`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("metrics", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    __table_args__ = (PrimaryKeyConstraint("key", "timestamp", "run_uuid", name="metric_pk"),)
+
+    def __repr__(self):
+        return f"<SqlMetric({self.key}, {self.value}, {self.timestamp})>"
+
+
+class SqlParam(Base):
+    __tablename__ = "params"
+
+    key = Column(String(250))
+    """
+    Param key: `String` (limit 250 characters). Part of *Primary Key* for ``params`` table.
+    """
+    value = Column(String(250), nullable=False)
+    """
+    Param value: `String` (limit 250 characters). Defined as *Non-null* in schema.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``params`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("params", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    __table_args__ = (PrimaryKeyConstraint("key", "run_uuid", name="param_pk"),)
+
+    def __repr__(self):
+        return f"<SqlParam({self.key}, {self.value})>"

venv/lib/python3.9/site-packages/mlflow/store/tracking/dbmodels/models.py

@@ -0,0 +1,755 @@
+import sqlalchemy as sa
+from sqlalchemy import (
+    BigInteger,
+    Boolean,
+    CheckConstraint,
+    Column,
+    ForeignKey,
+    Index,
+    Integer,
+    PrimaryKeyConstraint,
+    String,
+    UnicodeText,
+)
+from sqlalchemy.orm import backref, relationship
+
+from mlflow.entities import (
+    Dataset,
+    Experiment,
+    ExperimentTag,
+    InputTag,
+    Metric,
+    Param,
+    Run,
+    RunData,
+    RunInfo,
+    RunStatus,
+    RunTag,
+    SourceType,
+    TraceInfo,
+    ViewType,
+)
+from mlflow.entities.lifecycle_stage import LifecycleStage
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_status import TraceStatus
+from mlflow.store.db.base_sql_model import Base
+from mlflow.utils.mlflow_tags import _get_run_name_from_tags
+from mlflow.utils.time import get_current_time_millis
+
+SourceTypes = [
+    SourceType.to_string(SourceType.NOTEBOOK),
+    SourceType.to_string(SourceType.JOB),
+    SourceType.to_string(SourceType.LOCAL),
+    SourceType.to_string(SourceType.UNKNOWN),
+    SourceType.to_string(SourceType.PROJECT),
+]
+
+RunStatusTypes = [
+    RunStatus.to_string(RunStatus.SCHEDULED),
+    RunStatus.to_string(RunStatus.FAILED),
+    RunStatus.to_string(RunStatus.FINISHED),
+    RunStatus.to_string(RunStatus.RUNNING),
+    RunStatus.to_string(RunStatus.KILLED),
+]
+
+
+class SqlExperiment(Base):
+    """
+    DB model for :py:class:`mlflow.entities.Experiment`. These are recorded in ``experiment`` table.
+    """
+
+    __tablename__ = "experiments"
+
+    experiment_id = Column(Integer, autoincrement=True)
+    """
+    Experiment ID: `Integer`. *Primary Key* for ``experiment`` table.
+    """
+    name = Column(String(256), unique=True, nullable=False)
+    """
+    Experiment name: `String` (limit 256 characters). Defined as *Unique* and *Non null* in
+                     table schema.
+    """
+    artifact_location = Column(String(256), nullable=True)
+    """
+    Default artifact location for this experiment: `String` (limit 256 characters). Defined as
+                                                    *Non null* in table schema.
+    """
+    lifecycle_stage = Column(String(32), default=LifecycleStage.ACTIVE)
+    """
+    Lifecycle Stage of experiment: `String` (limit 32 characters).
+                                    Can be either ``active`` (default) or ``deleted``.
+    """
+    creation_time = Column(BigInteger(), default=get_current_time_millis)
+    """
+    Creation time of experiment: `BigInteger`.
+    """
+    last_update_time = Column(BigInteger(), default=get_current_time_millis)
+    """
+    Last Update time of experiment: `BigInteger`.
+    """
+
+    __table_args__ = (
+        CheckConstraint(
+            lifecycle_stage.in_(LifecycleStage.view_type_to_stages(ViewType.ALL)),
+            name="experiments_lifecycle_stage",
+        ),
+        PrimaryKeyConstraint("experiment_id", name="experiment_pk"),
+    )
+
+    def __repr__(self):
+        return f"<SqlExperiment ({self.experiment_id}, {self.name})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            :py:class:`mlflow.entities.Experiment`.
+        """
+        return Experiment(
+            experiment_id=str(self.experiment_id),
+            name=self.name,
+            artifact_location=self.artifact_location,
+            lifecycle_stage=self.lifecycle_stage,
+            tags=[t.to_mlflow_entity() for t in self.tags],
+            creation_time=self.creation_time,
+            last_update_time=self.last_update_time,
+        )
+
+
+class SqlRun(Base):
+    """
+    DB model for :py:class:`mlflow.entities.Run`. These are recorded in ``runs`` table.
+    """
+
+    __tablename__ = "runs"
+
+    run_uuid = Column(String(32), nullable=False)
+    """
+    Run UUID: `String` (limit 32 characters). *Primary Key* for ``runs`` table.
+    """
+    name = Column(String(250))
+    """
+    Run name: `String` (limit 250 characters).
+    """
+    source_type = Column(String(20), default=SourceType.to_string(SourceType.LOCAL))
+    """
+    Source Type: `String` (limit 20 characters). Can be one of ``NOTEBOOK``, ``JOB``, ``PROJECT``,
+                 ``LOCAL`` (default), or ``UNKNOWN``.
+    """
+    source_name = Column(String(500))
+    """
+    Name of source recording the run: `String` (limit 500 characters).
+    """
+    entry_point_name = Column(String(50))
+    """
+    Entry-point name that launched the run run: `String` (limit 50 characters).
+    """
+    user_id = Column(String(256), nullable=True, default=None)
+    """
+    User ID: `String` (limit 256 characters). Defaults to ``null``.
+    """
+    status = Column(String(20), default=RunStatus.to_string(RunStatus.SCHEDULED))
+    """
+    Run Status: `String` (limit 20 characters). Can be one of ``RUNNING``, ``SCHEDULED`` (default),
+                ``FINISHED``, ``FAILED``.
+    """
+    start_time = Column(BigInteger, default=get_current_time_millis)
+    """
+    Run start time: `BigInteger`. Defaults to current system time.
+    """
+    end_time = Column(BigInteger, nullable=True, default=None)
+    """
+    Run end time: `BigInteger`.
+    """
+    deleted_time = Column(BigInteger, nullable=True, default=None)
+    """
+    Run deleted time: `BigInteger`. Timestamp of when run is deleted, defaults to none.
+    """
+    source_version = Column(String(50))
+    """
+    Source version: `String` (limit 50 characters).
+    """
+    lifecycle_stage = Column(String(20), default=LifecycleStage.ACTIVE)
+    """
+    Lifecycle Stage of run: `String` (limit 32 characters).
+                            Can be either ``active`` (default) or ``deleted``.
+    """
+    artifact_uri = Column(String(200), default=None)
+    """
+    Default artifact location for this run: `String` (limit 200 characters).
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
+    """
+    Experiment ID to which this run belongs to: *Foreign Key* into ``experiment`` table.
+    """
+    experiment = relationship("SqlExperiment", backref=backref("runs", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`.
+    """
+
+    __table_args__ = (
+        CheckConstraint(source_type.in_(SourceTypes), name="source_type"),
+        CheckConstraint(status.in_(RunStatusTypes), name="status"),
+        CheckConstraint(
+            lifecycle_stage.in_(LifecycleStage.view_type_to_stages(ViewType.ALL)),
+            name="runs_lifecycle_stage",
+        ),
+        PrimaryKeyConstraint("run_uuid", name="run_pk"),
+    )
+
+    @staticmethod
+    def get_attribute_name(mlflow_attribute_name):
+        """
+        Resolves an MLflow attribute name to a `SqlRun` attribute name.
+        """
+        # Currently, MLflow Search attributes defined in `SearchUtils.VALID_SEARCH_ATTRIBUTE_KEYS`
+        # share the same names as their corresponding `SqlRun` attributes. Therefore, this function
+        # returns the same attribute name
+        return {"run_name": "name", "run_id": "run_uuid"}.get(
+            mlflow_attribute_name, mlflow_attribute_name
+        )
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Run: Description of the return value.
+        """
+        run_info = RunInfo(
+            run_uuid=self.run_uuid,
+            run_id=self.run_uuid,
+            run_name=self.name,
+            experiment_id=str(self.experiment_id),
+            user_id=self.user_id,
+            status=self.status,
+            start_time=self.start_time,
+            end_time=self.end_time,
+            lifecycle_stage=self.lifecycle_stage,
+            artifact_uri=self.artifact_uri,
+        )
+
+        tags = [t.to_mlflow_entity() for t in self.tags]
+        run_data = RunData(
+            metrics=[m.to_mlflow_entity() for m in self.latest_metrics],
+            params=[p.to_mlflow_entity() for p in self.params],
+            tags=tags,
+        )
+        if not run_info.run_name:
+            run_name = _get_run_name_from_tags(tags)
+            if run_name:
+                run_info._set_run_name(run_name)
+
+        return Run(run_info=run_info, run_data=run_data)
+
+
+class SqlExperimentTag(Base):
+    """
+    DB model for :py:class:`mlflow.entities.RunTag`.
+    These are recorded in ``experiment_tags`` table.
+    """
+
+    __tablename__ = "experiment_tags"
+
+    key = Column(String(250))
+    """
+    Tag key: `String` (limit 250 characters). *Primary Key* for ``tags`` table.
+    """
+    value = Column(String(5000), nullable=True)
+    """
+    Value associated with tag: `String` (limit 5000 characters). Could be *null*.
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"))
+    """
+    Experiment ID to which this tag belongs: *Foreign Key* into ``experiments`` table.
+    """
+    experiment = relationship("SqlExperiment", backref=backref("tags", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlExperiment`.
+    """
+
+    __table_args__ = (PrimaryKeyConstraint("key", "experiment_id", name="experiment_tag_pk"),)
+
+    def __repr__(self):
+        return f"<SqlExperimentTag({self.key}, {self.value})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.RunTag: Description of the return value.
+        """
+        return ExperimentTag(key=self.key, value=self.value)
+
+
+class SqlTag(Base):
+    """
+    DB model for :py:class:`mlflow.entities.RunTag`. These are recorded in ``tags`` table.
+    """
+
+    __tablename__ = "tags"
+    __table_args__ = (
+        PrimaryKeyConstraint("key", "run_uuid", name="tag_pk"),
+        Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+    )
+
+    key = Column(String(250))
+    """
+    Tag key: `String` (limit 250 characters). *Primary Key* for ``tags`` table.
+    """
+    value = Column(String(8000), nullable=True)
+    """
+    Value associated with tag: `String` (limit 8000 characters). Could be *null*.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this tag belongs to: *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("tags", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    def __repr__(self):
+        return f"<SqlRunTag({self.key}, {self.value})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            :py:class:`mlflow.entities.RunTag`.
+        """
+        return RunTag(key=self.key, value=self.value)
+
+
+class SqlMetric(Base):
+    __tablename__ = "metrics"
+    __table_args__ = (
+        PrimaryKeyConstraint(
+            "key", "timestamp", "step", "run_uuid", "value", "is_nan", name="metric_pk"
+        ),
+        Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+    )
+
+    key = Column(String(250))
+    """
+    Metric key: `String` (limit 250 characters). Part of *Primary Key* for ``metrics`` table.
+    """
+    value = Column(sa.types.Float(precision=53), nullable=False)
+    """
+    Metric value: `Float`. Defined as *Non-null* in schema.
+    """
+    timestamp = Column(BigInteger, default=get_current_time_millis)
+    """
+    Timestamp recorded for this metric entry: `BigInteger`. Part of *Primary Key* for
+                                               ``metrics`` table.
+    """
+    step = Column(BigInteger, default=0, nullable=False)
+    """
+    Step recorded for this metric entry: `BigInteger`.
+    """
+    is_nan = Column(Boolean(create_constraint=True), nullable=False, default=False)
+    """
+    True if the value is in fact NaN.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``metrics`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("metrics", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    def __repr__(self):
+        return f"<SqlMetric({self.key}, {self.value}, {self.timestamp}, {self.step})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Metric: Description of the return value.
+        """
+        return Metric(
+            key=self.key,
+            value=self.value if not self.is_nan else float("nan"),
+            timestamp=self.timestamp,
+            step=self.step,
+        )
+
+
+class SqlLatestMetric(Base):
+    __tablename__ = "latest_metrics"
+    __table_args__ = (
+        PrimaryKeyConstraint("key", "run_uuid", name="latest_metric_pk"),
+        Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+    )
+
+    key = Column(String(250))
+    """
+    Metric key: `String` (limit 250 characters). Part of *Primary Key* for ``latest_metrics`` table.
+    """
+    value = Column(sa.types.Float(precision=53), nullable=False)
+    """
+    Metric value: `Float`. Defined as *Non-null* in schema.
+    """
+    timestamp = Column(BigInteger, default=get_current_time_millis)
+    """
+    Timestamp recorded for this metric entry: `BigInteger`. Part of *Primary Key* for
+                                               ``latest_metrics`` table.
+    """
+    step = Column(BigInteger, default=0, nullable=False)
+    """
+    Step recorded for this metric entry: `BigInteger`.
+    """
+    is_nan = Column(Boolean(create_constraint=True), nullable=False, default=False)
+    """
+    True if the value is in fact NaN.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``latest_metrics`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("latest_metrics", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    def __repr__(self):
+        return f"<SqlLatestMetric({self.key}, {self.value}, {self.timestamp}, {self.step})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Metric: Description of the return value.
+        """
+        return Metric(
+            key=self.key,
+            value=self.value if not self.is_nan else float("nan"),
+            timestamp=self.timestamp,
+            step=self.step,
+        )
+
+
+class SqlParam(Base):
+    __tablename__ = "params"
+    __table_args__ = (
+        PrimaryKeyConstraint("key", "run_uuid", name="param_pk"),
+        Index(f"index_{__tablename__}_run_uuid", "run_uuid"),
+    )
+
+    key = Column(String(250))
+    """
+    Param key: `String` (limit 250 characters). Part of *Primary Key* for ``params`` table.
+    """
+    value = Column(String(8000), nullable=False)
+    """
+    Param value: `String` (limit 8000 characters). Defined as *Non-null* in schema.
+    """
+    run_uuid = Column(String(32), ForeignKey("runs.run_uuid"))
+    """
+    Run UUID to which this metric belongs to: Part of *Primary Key* for ``params`` table.
+                                              *Foreign Key* into ``runs`` table.
+    """
+    run = relationship("SqlRun", backref=backref("params", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with :py:class:`mlflow.store.dbmodels.models.SqlRun`.
+    """
+
+    def __repr__(self):
+        return f"<SqlParam({self.key}, {self.value})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Param: Description of the return value.
+        """
+        return Param(key=self.key, value=self.value)
+
+
+class SqlDataset(Base):
+    __tablename__ = "datasets"
+    __table_args__ = (
+        PrimaryKeyConstraint("experiment_id", "name", "digest", name="dataset_pk"),
+        Index(f"index_{__tablename__}_dataset_uuid", "dataset_uuid"),
+        Index(
+            f"index_{__tablename__}_experiment_id_dataset_source_type",
+            "experiment_id",
+            "dataset_source_type",
+        ),
+    )
+
+    dataset_uuid = Column(String(36), nullable=False)
+    """
+    Dataset UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``datasets`` table.
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id", ondelete="CASCADE"))
+    """
+    Experiment ID to which this dataset belongs: *Foreign Key* into ``experiments`` table.
+    """
+    name = Column(String(500), nullable=False)
+    """
+    Param name: `String` (limit 500 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``datasets`` table.
+    """
+    digest = Column(String(36), nullable=False)
+    """
+    Param digest: `String` (limit 500 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``datasets`` table.
+    """
+    dataset_source_type = Column(String(36), nullable=False)
+    """
+    Param dataset_source_type: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    """
+    dataset_source = Column(UnicodeText, nullable=False)
+    """
+    Param dataset_source: `UnicodeText`. Defined as *Non-null* in schema.
+    """
+    dataset_schema = Column(UnicodeText, nullable=True)
+    """
+    Param dataset_schema: `UnicodeText`.
+    """
+    dataset_profile = Column(UnicodeText, nullable=True)
+    """
+    Param dataset_profile: `UnicodeText`.
+    """
+
+    def __repr__(self):
+        return "<SqlDataset ({}, {}, {}, {}, {}, {}, {}, {})>".format(
+            self.dataset_uuid,
+            self.experiment_id,
+            self.name,
+            self.digest,
+            self.dataset_source_type,
+            self.dataset_source,
+            self.dataset_schema,
+            self.dataset_profile,
+        )
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.Dataset.
+        """
+        return Dataset(
+            name=self.name,
+            digest=self.digest,
+            source_type=self.dataset_source_type,
+            source=self.dataset_source,
+            schema=self.dataset_schema,
+            profile=self.dataset_profile,
+        )
+
+
+class SqlInput(Base):
+    __tablename__ = "inputs"
+    __table_args__ = (
+        PrimaryKeyConstraint(
+            "source_type", "source_id", "destination_type", "destination_id", name="inputs_pk"
+        ),
+        Index(f"index_{__tablename__}_input_uuid", "input_uuid"),
+        Index(
+            f"index_{__tablename__}_destination_type_destination_id_source_type",
+            "destination_type",
+            "destination_id",
+            "source_type",
+        ),
+    )
+
+    input_uuid = Column(String(36), nullable=False)
+    """
+    Input UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    """
+    source_type = Column(String(36), nullable=False)
+    """
+    Source type: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``inputs`` table.
+    """
+    source_id = Column(String(36), nullable=False)
+    """
+    Source Id: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``inputs`` table.
+    """
+    destination_type = Column(String(36), nullable=False)
+    """
+    Destination type: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``inputs`` table.
+    """
+    destination_id = Column(String(36), nullable=False)
+    """
+    Destination Id: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``inputs`` table.
+    """
+
+    def __repr__(self):
+        return "<SqlInput ({}, {}, {}, {}, {})>".format(
+            self.input_uuid,
+            self.source_type,
+            self.source_id,
+            self.destination_type,
+            self.destination_id,
+        )
+
+
+class SqlInputTag(Base):
+    __tablename__ = "input_tags"
+    __table_args__ = (PrimaryKeyConstraint("input_uuid", "name", name="input_tags_pk"),)
+
+    input_uuid = Column(String(36), ForeignKey("inputs.input_uuid"), nullable=False)
+    """
+    Input UUID: `String` (limit 36 characters). Defined as *Non-null* in schema.
+    *Foreign Key* into ``inputs`` table. Part of *Primary Key* for ``input_tags`` table.
+    """
+    name = Column(String(255), nullable=False)
+    """
+    Param name: `String` (limit 255 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``input_tags`` table.
+    """
+    value = Column(String(500), nullable=False)
+    """
+    Param value: `String` (limit 500 characters). Defined as *Non-null* in schema.
+    Part of *Primary Key* for ``input_tags`` table.
+    """
+
+    def __repr__(self):
+        return f"<SqlInputTag ({self.input_uuid}, {self.name}, {self.value})>"
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            mlflow.entities.InputTag: Description of the return value.
+        """
+        return InputTag(key=self.name, value=self.value)
+
+
+#######################################################################################
+# Below are Tracing models. We may refactor them to be in a separate module in the future.
+#######################################################################################
+
+
+class SqlTraceInfo(Base):
+    __tablename__ = "trace_info"
+
+    request_id = Column(String(50), nullable=False)
+    """
+    Request ID: `String` (limit 50 characters). *Primary Key* for ``trace_info`` table.
+    """
+    experiment_id = Column(Integer, ForeignKey("experiments.experiment_id"), nullable=False)
+    """
+    Experiment ID to which this trace belongs: *Foreign Key* into ``experiments`` table.
+    """
+    timestamp_ms = Column(BigInteger, nullable=False)
+    """
+    Start time of the trace, in milliseconds.
+    """
+    execution_time_ms = Column(BigInteger, nullable=True)
+    """
+    Duration of the trace, in milliseconds. Could be *null* if the trace is still in progress
+    or not ended correctly for some reason.
+    """
+    status = Column(String(50), nullable=False)
+    """
+    Status of the trace. The values are defined in
+    :py:class:`mlflow.entities.trace_status.TraceStatus` enum but we don't enforce
+    constraint at DB level.
+    """
+
+    __table_args__ = (
+        PrimaryKeyConstraint("request_id", name="trace_info_pk"),
+        # The most frequent query will be get all traces in an experiment sorted by timestamp desc,
+        # which is the default view in the UI. Also every search query should have experiment_id(s)
+        # in the where clause.
+        Index(f"index_{__tablename__}_experiment_id_timestamp_ms", "experiment_id", "timestamp_ms"),
+    )
+
+    def to_mlflow_entity(self):
+        """
+        Convert DB model to corresponding MLflow entity.
+
+        Returns:
+            :py:class:`mlflow.entities.TraceInfo` object.
+        """
+        return TraceInfo(
+            request_id=self.request_id,
+            experiment_id=str(self.experiment_id),
+            timestamp_ms=self.timestamp_ms,
+            execution_time_ms=self.execution_time_ms,
+            status=TraceStatus(self.status),
+            tags={t.key: t.value for t in self.tags},
+            request_metadata={m.key: m.value for m in self.request_metadata},
+        )
+
+
+class SqlTraceTag(Base):
+    __tablename__ = "trace_tags"
+
+    key = Column(String(250))
+    """
+    Tag key: `String` (limit 250 characters).
+    """
+    value = Column(String(8000), nullable=True)
+    """
+    Value associated with tag: `String` (limit 250 characters). Could be *null*.
+    """
+    request_id = Column(
+        String(50), ForeignKey("trace_info.request_id", ondelete="CASCADE"), nullable=False
+    )
+    """
+    Request ID to which this tag belongs: *Foreign Key* into ``trace_info`` table.
+    """
+    trace_info = relationship("SqlTraceInfo", backref=backref("tags", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with
+    :py:class:`mlflow.store.dbmodels.models.SqlTraceInfo`.
+    """
+
+    # Key is unique within a request_id
+    __table_args__ = (
+        PrimaryKeyConstraint("request_id", "key", name="trace_tag_pk"),
+        Index(f"index_{__tablename__}_request_id"),
+    )
+
+
+class SqlTraceRequestMetadata(Base):
+    __tablename__ = "trace_request_metadata"
+
+    key = Column(String(250))
+    """
+    Metadata key: `String` (limit 250 characters).
+    """
+    value = Column(String(8000), nullable=True)
+    """
+    Value associated with metadata: `String` (limit 250 characters). Could be *null*.
+    """
+    request_id = Column(
+        String(50), ForeignKey("trace_info.request_id", ondelete="CASCADE"), nullable=False
+    )
+    """
+    Request ID to which this metadata belongs: *Foreign Key* into ``trace_info`` table.
+    """
+    trace_info = relationship("SqlTraceInfo", backref=backref("request_metadata", cascade="all"))
+    """
+    SQLAlchemy relationship (many:one) with
+    :py:class:`mlflow.store.dbmodels.models.SqlTraceInfo`.
+    """
+
+    # Key is unique within a request_id
+    __table_args__ = (
+        PrimaryKeyConstraint("request_id", "key", name="trace_request_metadata_pk"),
+        Index(f"index_{__tablename__}_request_id"),
+    )