Add labels to Experiment (#11)

* Add Roadmap

Signed-off-by: kerthcet <kerthcet@gmail.com>

* Add labels

Signed-off-by: kerthcet <kerthcet@gmail.com>

---------

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

Makefile

@@ -14,4 +14,9 @@ launch:
 
 .PHONY: test
 test:
-	$(POETRY) run pytest
+	$(POETRY) run pytest
+
+.PHONY: format
+format:
+	ruff format .
+	ruff check --fix .

alphatrion/consts.py

@@ -1 +1 @@
-METADATA_DB_URL = "METADATA_DB_URL"
+METADATA_DB_URL = "METADATA_DB_URL"

alphatrion/experiment/base.py

@@ -1,29 +1,40 @@
-
-
-from abc import ABC
+from abc import ABC, abstractmethod
 
 from alphatrion.runtime.runtime import Runtime
 
 
 class Experiment(ABC):
     """Base class for all experiments."""
+
     def __init__(self, runtime: Runtime):
         self._runtime = runtime
 
-    def create(self, name: str, description: str | None = None, meta: dict | None = None):
+    @abstractmethod
+    def create(
+        self, name: str, description: str | None = None, meta: dict | None = None, labels: dict | None = None
+    ):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
     def delete(self, exp_id: int):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
     def get(self, exp_id: int):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
+    def update_labels(self, exp_id: int, labels: dict):
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    @abstractmethod
     def start(self, exp_id: int):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
     def stop(self, exp_id: int, status: str = "finished"):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
     def status(self, exp_id: int) -> str:
         raise NotImplementedError("Subclasses must implement this method.")

alphatrion/experiment/custom_exp.py

@@ -1,28 +1,47 @@
-from alphatrion.metadata.sql_models import ExperimentStatus
+from datetime import datetime
+
 from alphatrion.experiment.base import Experiment
+from alphatrion.metadata.sql_models import COMPLETED_STATUS, ExperimentStatus
 from alphatrion.runtime.runtime import Runtime
 
 
 class CustomExperiment(Experiment):
     def __init__(self, runtime: Runtime):
         super().__init__(runtime)
 
-    def create(self, name: str, description: str | None = None, meta: dict | None = None):
-        self._runtime._metadb.create_exp(name=name, description=description, project_id=self._runtime._project_id, meta=meta)
+    def create(
+        self, name: str, description: str | None = None, meta: dict | None = None, labels: dict | None = None
+    ):
+        self._runtime._metadb.create_exp(
+            name=name,
+            description=description,
+            project_id=self._runtime._project_id,
+            meta=meta,
+            labels=labels,
+        )
 
     def delete(self, exp_id: int):
         self._runtime._metadb.delete_exp(exp_id=exp_id)
 
     def get(self, exp_id: int):
         return self._runtime._metadb.get_exp(exp_id=exp_id)
 
+    # Please provide all the labels to update, or it will overwrite the existing labels.
+    def update_labels(self, exp_id: int, labels: dict):
+        self._runtime._metadb.update_exp(exp_id=exp_id, labels=labels)
+
     # start for experiment usually means update the status to running.
     def start(self, exp_id: int):
         self._runtime._metadb.update_exp(exp_id=exp_id, status=ExperimentStatus.RUNNING)
 
     # stop for experiment usually means update the status to finished or failed.
     def stop(self, exp_id: int, status: ExperimentStatus = ExperimentStatus.FINISHED):
-        self._runtime._metadb.update_exp(exp_id=exp_id, status=status)
+        exp = self._runtime._metadb.get_exp(exp_id=exp_id)
+        if exp is not None and exp.status not in COMPLETED_STATUS:
+            duration = (datetime.now() - exp.created_at).total_seconds()
+            self._runtime._metadb.update_exp(
+                exp_id=exp_id, status=status, duration=duration
+            )
 
     def status(self, exp_id: int) -> ExperimentStatus:
         exp = self._runtime._metadb.get_exp(exp_id=exp_id)

alphatrion/main.py

@@ -3,4 +3,4 @@
 load_dotenv()
 
 if __name__ == "__main__":
-    print("Hello, AlphaTrion!")
+    print("Hello, AlphaTrion!")

alphatrion/metadata/base.py

@@ -1,21 +1,27 @@
-from abc import ABC
+from abc import ABC, abstractmethod
+
 
 class MetaStore(ABC):
     """Base class for all metadata storage backends."""
-    def __init__(self):
-        pass
 
-    def create_exp(self, name: str, project_id: str, description: str | None, meta: dict | None):
+    @abstractmethod
+    def create_exp(
+        self, name: str, project_id: str, description: str | None, meta: dict | None
+    ):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
     def delete_exp(self, exp_id: int):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
     def update_exp(self, exp_id: int, **kwargs):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
     def get_exp(self, exp_id: int):
         raise NotImplementedError("Subclasses must implement this method.")
 
+    @abstractmethod
     def list_exps(self, project_id: str, page: int, page_size: int):
         raise NotImplementedError("Subclasses must implement this method.")

alphatrion/metadata/sql.py

@@ -1,26 +1,36 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
-from alphatrion.metadata.sql_models import Base, Experiment
 from alphatrion.metadata.base import MetaStore
+from alphatrion.metadata.sql_models import Base, Experiment
 
 
 # SQL-like metadata implementation, it could be SQLite, PostgreSQL, MySQL, etc.
 class SQLStore(MetaStore):
     def __init__(self, db_url: str, init_tables: bool = False):
-        super().__init__()
-
         self._engine = create_engine(db_url)
         self._session = sessionmaker(bind=self._engine)
         if init_tables:
             # create tables if not exist, will not affect existing tables.
             # In production, use migrations instead.
             Base.metadata.create_all(self._engine)
 
-
-    def create_exp(self, name: str, project_id: str, description: str | None, meta: dict | None):
+    def create_exp(
+        self,
+        name: str,
+        project_id: str,
+        description: str | None,
+        meta: dict | None,
+        labels: dict | None = None,
+    ):
         session = self._session()
-        new_exp = Experiment(name=name, description=description, project_id=project_id, meta=meta)
+        new_exp = Experiment(
+            name=name,
+            description=description,
+            project_id=project_id,
+            meta=meta,
+            labels=labels,
+        )
         session.add(new_exp)
         session.commit()
         session.close()
@@ -34,6 +44,7 @@ def delete_exp(self, exp_id: int):
             session.commit()
         session.close()
 
+    # We don't support append-only update, the complete fields should be provided.
     def update_exp(self, exp_id: int, **kwargs):
         session = self._session()
         exp = session.query(Experiment).filter(Experiment.id == exp_id).first()
@@ -46,13 +57,23 @@ def update_exp(self, exp_id: int, **kwargs):
     # get_exp will ignore the deleted experiments.
     def get_exp(self, exp_id: int) -> Experiment | None:
         session = self._session()
-        exp = session.query(Experiment).filter(Experiment.id == exp_id, Experiment.is_del == 0).first()
+        exp = (
+            session.query(Experiment)
+            .filter(Experiment.id == exp_id, Experiment.is_del == 0)
+            .first()
+        )
         session.close()
         return exp
 
     # paginate the experiments in case of too many experiments.
     def list_exps(self, project_id: str, page: int, page_size: int) -> list[Experiment]:
         session = self._session()
-        exps = session.query(Experiment).filter(Experiment.project_id == project_id).offset(page * page_size).limit(page_size).all()
+        exps = (
+            session.query(Experiment)
+            .filter(Experiment.project_id == project_id)
+            .offset(page * page_size)
+            .limit(page_size)
+            .all()
+        )
         session.close()
         return exps

alphatrion/metadata/sql_models.py

@@ -1,7 +1,7 @@
-from datetime import datetime, timezone
 import enum
+from datetime import UTC, datetime
 
-from sqlalchemy import Column, Integer, String, Enum, DateTime, JSON
+from sqlalchemy import JSON, Column, DateTime, Enum, Integer, String
 from sqlalchemy.orm import declarative_base
 
 Base = declarative_base()
@@ -14,6 +14,9 @@ class ExperimentStatus(enum.Enum):
     FAILED = "failed"
 
 
+COMPLETED_STATUS = [ExperimentStatus.FINISHED, ExperimentStatus.FAILED]
+
+
 # Define the Experiment model for SQLAlchemy
 class Experiment(Base):
     __tablename__ = "experiments"
@@ -22,9 +25,15 @@ class Experiment(Base):
     name = Column(String, nullable=False, unique=True)
     description = Column(String, nullable=True)
     project_id = Column(String, nullable=False)
-    status = Column(Enum(ExperimentStatus), nullable=False, default=ExperimentStatus.PENDING)
+    status = Column(
+        Enum(ExperimentStatus), nullable=False, default=ExperimentStatus.PENDING
+    )
     meta = Column(JSON, nullable=True, comment="Additional metadata for the experiment")
+    labels = Column(JSON, nullable=True, comment="Labels for the experiment")
+    duration = Column(Integer, default=0, comment="Duration in seconds")
 
-    created_at = Column(DateTime(timezone=True), default=datetime.now(timezone.utc))
-    updated_at = Column(DateTime(timezone=True), default=datetime.now(timezone.utc), onupdate=datetime.now(timezone.utc))
+    created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+    updated_at = Column(
+        DateTime(timezone=True), default=datetime.now(UTC), onupdate=datetime.now(UTC)
+    )
     is_del = Column(Integer, default=0, comment="0 for not deleted, 1 for deleted")

poetry.lock

@@ -13,7 +13,31 @@ dependencies = [
 
 [tool.poetry.group.dev.dependencies]
 pytest = ">=8.4.2,<9.0.0"
+ruff = "^0.12.12"
 
 [build-system]
 requires = ["poetry-core>=2.0.0,<3.0.0"]
 build-backend = "poetry.core.masonry.api"
+
+# Configuration for ruff linter and formatter
+
+[tool.ruff]
+line-length = 88
+target-version = "py312"
+lint.select = [
+    "E",  # pycodestyle,
+    "F",  # Pyflakes
+    "UP",  # pyupgrade
+    "B",  # flake8-bugbear
+    "SIM",  # flake8-simplify
+    "I",  # isort
+    "PL",  # PyLint,
+]
+exclude = ["venv", "migrations", "__pycache__"]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["PLR2004"]