Add project_id to all models (#46)

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

Author: kerthcet

alphatrion/experiment/craft_exp.py

@@ -1,4 +1,3 @@
-
 from alphatrion.experiment.base import Experiment
 from alphatrion.trial.trial import Trial, TrialConfig
 

alphatrion/log/log.py

@@ -68,6 +68,7 @@ async def log_metrics(metrics: dict[str, float]):
         runtime._metadb.create_metric(
             key=key,
             value=value,
+            project_id=runtime._project_id,
             trial_id=trial_id,
             step=step,
         )

alphatrion/metadata/sql.py

@@ -9,6 +9,7 @@
     Experiment,
     Metrics,
     Model,
+    Project,
     Run,
     Trial,
     TrialStatus,
@@ -178,7 +179,8 @@ def delete_model(self, model_id: int):
 
     def create_trial(
         self,
-        exp_id: int,
+        project_id: uuid.UUID,
+        exp_id: uuid.UUID,
         name: str,
         description: str | None = None,
         meta: dict | None = None,
@@ -187,6 +189,7 @@ def create_trial(
     ) -> uuid.UUID:
         session = self._session()
         new_trial = Trial(
+            project_id=project_id,
             experiment_id=exp_id,
             name=name,
             description=description,
@@ -217,9 +220,10 @@ def update_trial(self, trial_id: uuid.UUID, **kwargs):
             session.commit()
         session.close()
 
-    def create_run(self, trial_id: uuid.UUID) -> uuid.UUID:
+    def create_run(self, project_id: uuid.UUID, trial_id: uuid.UUID) -> uuid.UUID:
         session = self._session()
         new_run = Run(
+            project_id=project_id,
             trial_id=trial_id,
         )
         session.add(new_run)
@@ -228,9 +232,17 @@ def create_run(self, trial_id: uuid.UUID) -> uuid.UUID:
         session.close()
         return run_id
 
-    def create_metric(self, trial_id: uuid.UUID, key: str, value: float, step: int):
+    def create_metric(
+        self,
+        project_id: uuid.UUID,
+        trial_id: uuid.UUID,
+        key: str,
+        value: float,
+        step: int,
+    ):
         session = self._session()
         new_metric = Metrics(
+            project_id=project_id,
             trial_id=trial_id,
             key=key,
             value=value,
@@ -245,3 +257,13 @@ def list_metrics(self, trial_id: uuid.UUID) -> list[Metrics]:
         metrics = session.query(Metrics).filter(Metrics.trial_id == trial_id).all()
         session.close()
         return metrics
+
+    def get_project(self, project_id: str) -> Project | None:
+        session = self._session()
+        project = (
+            session.query(Project)
+            .filter(Project.uuid == project_id, Project.is_del == 0)
+            .first()
+        )
+        session.close()
+        return project

alphatrion/metadata/sql_models.py

@@ -19,14 +19,28 @@ class TrialStatus(enum.Enum):
 COMPLETED_STATUS = [TrialStatus.FINISHED, TrialStatus.FAILED]
 
 
+class Project(Base):
+    __tablename__ = "projects"
+
+    uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    name = Column(String, nullable=False)
+    description = Column(String, nullable=True)
+
+    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")
+
+
 # Define the Experiment model for SQLAlchemy
 class Experiment(Base):
     __tablename__ = "experiments"
 
     uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     name = Column(String, nullable=False)
     description = Column(String, nullable=True)
-    project_id = Column(String, nullable=False)
+    project_id = Column(UUID(as_uuid=True), nullable=False)
     meta = Column(JSON, nullable=True, comment="Additional metadata for the experiment")
 
     created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
@@ -40,6 +54,7 @@ class Trial(Base):
     __tablename__ = "trials"
 
     uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    project_id = Column(UUID(as_uuid=True), nullable=False)
     experiment_id = Column(UUID(as_uuid=True), nullable=False)
     name = Column(String, nullable=False)
     description = Column(String, nullable=True)
@@ -63,8 +78,8 @@ class Run(Base):
     __tablename__ = "runs"
 
     uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    project_id = Column(UUID(as_uuid=True), nullable=False)
     trial_id = Column(UUID(as_uuid=True), nullable=False)
-    # artifact_path = Column(String, nullable=False)
 
     created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
     updated_at = Column(
@@ -77,10 +92,10 @@ class Model(Base):
 
     uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     name = Column(String, nullable=False, unique=True)
-    version = Column(String, nullable=False)
     description = Column(String, nullable=True)
+    project_id = Column(UUID(as_uuid=True), nullable=False)
+    version = Column(String, nullable=False)
     meta = Column(JSON, nullable=True, comment="Additional metadata for the model")
-    project_id = Column(String, nullable=False)
 
     created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
     updated_at = Column(
@@ -95,6 +110,7 @@ class Metrics(Base):
     uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
     key = Column(String, nullable=False)
     value = Column(Float, nullable=False)
+    project_id = Column(UUID(as_uuid=True), nullable=False)
     trial_id = Column(UUID(as_uuid=True), nullable=False)
     step = Column(Integer, nullable=False, default=0)
     created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))

alphatrion/run/run.py

@@ -13,4 +13,6 @@ def id(self) -> uuid.UUID:
         return self._id
 
     def _start(self):
-        self._id = self._runtime._metadb.create_run(trial_id=self._trial_id)
+        self._id = self._runtime._metadb.create_run(
+            project_id=self._runtime._project_id, trial_id=self._trial_id
+        )

alphatrion/runtime/runtime.py

@@ -1,5 +1,6 @@
 # ruff: noqa: PLW0603
 import os
+import uuid
 
 from alphatrion import consts
 from alphatrion.artifact.artifact import Artifact
@@ -9,7 +10,7 @@
 
 
 def init(
-    project_id: str = "alphatrion",
+    project_id: uuid.UUID,
     artifact_insecure: bool = False,
     init_tables: bool = False,
 ):

alphatrion/trial/trial.py

@@ -240,6 +240,7 @@ def _start(
         params: dict | None = None,
     ):
         self._id = self._runtime._metadb.create_trial(
+            project_id=self._runtime._project_id,
             exp_id=self._exp_id,
             name=name,
             description=description,

tests/integration/test_artifact.py

@@ -2,6 +2,7 @@
 
 import os
 import tempfile
+import uuid
 
 import pytest
 
@@ -10,14 +11,14 @@
 
 @pytest.fixture
 def artifact():
-    init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
     artifact = global_runtime()._artifact
 
     yield artifact
 
 
 def test_push_with_files(artifact):
-    init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     with tempfile.TemporaryDirectory() as tmpdir:
         os.chdir(tmpdir)
@@ -40,7 +41,7 @@ def test_push_with_files(artifact):
 
 
 def test_push_with_folder(artifact):
-    init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     with tempfile.TemporaryDirectory() as tmpdir:
         os.chdir(tmpdir)

tests/integration/test_log.py

@@ -2,6 +2,7 @@
 import os
 import tempfile
 import time
+import uuid
 from datetime import datetime, timedelta
 
 import pytest
@@ -13,7 +14,7 @@
 
 @pytest.mark.asyncio
 async def test_log_artifact():
-    alpha.init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    alpha.init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     async with alpha.CraftExperiment.start(
         name="log_artifact_exp",
@@ -65,7 +66,7 @@ async def test_log_artifact():
 
 @pytest.mark.asyncio
 async def test_log_params():
-    alpha.init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    alpha.init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     async with alpha.CraftExperiment.start(name="log_params_exp") as exp:
         trial = exp.start_trial(name="first-trial", params={"param1": 0.1})
@@ -92,7 +93,7 @@ async def test_log_params():
 
 @pytest.mark.asyncio
 async def test_log_metrics():
-    alpha.init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    alpha.init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     async with alpha.CraftExperiment.start(name="log_metrics_exp") as exp:
         trial = exp.start_trial(name="first-trial", params={"param1": 0.1})
@@ -128,7 +129,7 @@ async def test_log_metrics():
 
 @pytest.mark.asyncio
 async def test_log_metrics_with_save_on_max():
-    alpha.init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    alpha.init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     async with alpha.CraftExperiment.start(
         name="log_metrics_with_save_on_max",
@@ -182,7 +183,7 @@ async def test_log_metrics_with_save_on_max():
 
 @pytest.mark.asyncio
 async def test_log_metrics_with_save_on_min():
-    alpha.init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    alpha.init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     async with alpha.CraftExperiment.start(
         name="log_metrics_with_save_on_min",
@@ -236,7 +237,7 @@ async def test_log_metrics_with_save_on_min():
 
 @pytest.mark.asyncio
 async def test_log_metrics_with_early_stopping():
-    alpha.init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    alpha.init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     async def fake_work(value: float):
         await alpha.log_metrics({"accuracy": value})
@@ -273,7 +274,7 @@ async def fake_sleep(value: float):
 
 @pytest.mark.asyncio
 async def test_log_metrics_with_early_stopping_never_triggered():
-    alpha.init(project_id="test_project", artifact_insecure=True, init_tables=True)
+    alpha.init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
 
     async def fake_work(value: float):
         await alpha.log_metrics({"accuracy": value})

tests/unit/artifact/test_artifact.py

@@ -1,14 +1,16 @@
 # Test the Artifact class
 
 
+import uuid
+
 import pytest
 
 from alphatrion.runtime.runtime import global_runtime, init
 
 
 @pytest.fixture
 def artifact():
-    init(project_id="test_project", artifact_insecure=True)
+    init(project_id=uuid.uuid4(), artifact_insecure=True)
     artifact = global_runtime()._artifact
     yield artifact