Add seed script (#54)

* Add seed script

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

* Update Readme.md

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

* fix lint

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

---------

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

Author: kerthcet

Makefile

@@ -42,3 +42,7 @@ test-integration: lint
 	'
 .PHONY: test-all
 test-all: test test-integration
+
+.PHONY: seed
+seed:
+	python hack/seed.py

README.md

@@ -16,6 +16,7 @@ Open, modular framework to build GenAI applications.
 
 ## Concepts
 
+- **Project**: A Project is a namespace-level abstraction that isolates experiments from different users or teams.
 - **Experiment**: An Experiment is a high-level abstraction for organizing and managing a series of related trials. It serves as a way to group together multiple trials that share a common goal or objective.
 - **Trial**: A Trial represents a single attempt or multiple iterations within an experiment. It encapsulates the configuration, execution, and results of a specific set of runs.
 - **Run**: A Run is an execution of a specific configuration within a trial. It represents a real iteration of the trial.
@@ -67,7 +68,6 @@ Dashboard is coming soon! Meanwhile, you can query the results directly from the
 make down
 ```
 
-
 ## Contributing
 
 We welcome contributions! Please refer to [developer.md](./site/docs/development.md) for more information on how to set up your development environment and contribute to the project.

alphatrion/metadata/sql.py

@@ -7,7 +7,7 @@
 from alphatrion.metadata.sql_models import (
     Base,
     Experiment,
-    Metrics,
+    Metric,
     Model,
     Project,
     Run,
@@ -267,7 +267,7 @@ def create_metric(
         step: int,
     ):
         session = self._session()
-        new_metric = Metrics(
+        new_metric = Metric(
             project_id=project_id,
             trial_id=trial_id,
             key=key,
@@ -278,8 +278,8 @@ def create_metric(
         session.commit()
         session.close()
 
-    def list_metrics(self, trial_id: uuid.UUID) -> list[Metrics]:
+    def list_metrics(self, trial_id: uuid.UUID) -> list[Metric]:
         session = self._session()
-        metrics = session.query(Metrics).filter(Metrics.trial_id == trial_id).all()
+        metrics = session.query(Metric).filter(Metric.trial_id == trial_id).all()
         session.close()
         return metrics

alphatrion/metadata/sql_models.py

@@ -59,7 +59,6 @@ class Trial(Base):
     name = Column(String, nullable=False)
     description = Column(String, nullable=True)
     meta = Column(JSON, nullable=True, comment="Additional metadata for the trial")
-    duration = Column(Integer, default=0, comment="Duration in seconds")
     params = Column(JSON, nullable=True, comment="Parameters for the experiment")
     status = Column(
         Enum(TrialStatus),
@@ -104,7 +103,7 @@ class Model(Base):
     is_del = Column(Integer, default=0, comment="0 for not deleted, 1 for deleted")
 
 
-class Metrics(Base):
+class Metric(Base):
     __tablename__ = "metrics"
 
     uuid = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)

alphatrion/runtime/runtime.py

@@ -17,8 +17,8 @@ def init(
     """
     Initialize the AlphaTrion runtime environment.
 
-    :param project_id: the project ID to initialize the environment for,
-                       for testing purpose, you can use a random UUID.
+    :param project_id: the project ID to initialize the environment.
+                       For testing purpose, you can use a random UUID.
     :param artifact_insecure: whether to use insecure connection to the
         artifact registry
     """

alphatrion/tracing/__init__.py

@@ -0,0 +1,144 @@
+#!/usr/bin/env python3
+
+import os
+import random
+import uuid
+from datetime import datetime
+from decimal import Decimal
+
+from dotenv import load_dotenv
+from faker import Faker
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+from alphatrion import consts
+from alphatrion.metadata.sql_models import (
+    Base,
+    Experiment,
+    Metric,
+    Project,
+    Run,
+    Trial,
+    TrialStatus,
+)
+
+load_dotenv()
+
+DATABASE_URL = os.getenv(consts.METADATA_DB_URL)
+
+engine = create_engine(DATABASE_URL, echo=False)
+SessionLocal = sessionmaker(bind=engine)
+session = SessionLocal()
+
+fake = Faker()
+
+
+def make_json_serializable(obj):
+    if isinstance(obj, dict):
+        return {k: make_json_serializable(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [make_json_serializable(v) for v in obj]
+    elif isinstance(obj, Decimal):
+        return float(obj)
+    elif isinstance(obj, datetime):
+        return obj.isoformat()
+    else:
+        return obj
+
+
+def generate_project() -> Project:
+    return Project(
+        uuid=uuid.uuid4(),
+        name=fake.bs().title(),
+        description=fake.catch_phrase(),
+    )
+
+
+def generate_experiment(projects: list[Project]) -> Experiment:
+    return Experiment(
+        name=fake.bs().title(),
+        description=fake.catch_phrase(),
+        meta=make_json_serializable(
+            fake.pydict(nb_elements=3, variable_nb_elements=True)
+        ),
+        project_id=random.choice(projects).uuid,
+    )
+
+
+def generate_trial(exps: list[Experiment]) -> Trial:
+    exp = random.choice(exps)
+    return Trial(
+        project_id=exp.project_id,
+        experiment_id=exp.uuid,
+        name=fake.bs().title(),
+        description=fake.catch_phrase(),
+        meta=make_json_serializable(
+            fake.pydict(nb_elements=3, variable_nb_elements=True)
+        ),
+        params=make_json_serializable(
+            fake.pydict(nb_elements=3, variable_nb_elements=True)
+        ),
+        status=random.choice(list(TrialStatus)),
+    )
+
+
+def generate_run(trials: list[Trial]) -> Run:
+    trial = random.choice(trials)
+    return Run(
+        project_id=trial.project_id,
+        trial_id=trial.uuid,
+    )
+
+
+def generate_metric(runs: list[Run]) -> Metric:
+    run = random.choice(runs)
+    return Metric(
+        project_id=run.project_id,
+        trial_id=run.trial_id,
+        key=random.choice(["accuracy", "loss", "precision", "fitness"]),
+        value=random.uniform(0, 1),
+        step=random.randint(1, 1000),
+    )
+
+
+def seed_all(
+    num_projects: int,
+    num_exps_per_project: int,
+    num_trials_per_exp: int,
+    num_runs_per_trial: int,
+    num_metrics_per_run: int,
+):
+    Base.metadata.create_all(bind=engine)
+
+    print("🌱 generating seeds ...")
+    projects = [generate_project() for _ in range(num_projects)]
+    session.add_all(projects)
+    session.commit()
+
+    experiments = [generate_experiment(projects) for _ in range(num_exps_per_project)]
+    session.add_all(experiments)
+    session.commit()
+
+    trials = [generate_trial([exp]) for exp in experiments]
+    session.add_all(trials)
+    session.commit()
+
+    runs = [generate_run(trials) for _ in range(num_runs_per_trial)]
+    session.add_all(runs)
+    session.commit()
+
+    metrics = [generate_metric(runs) for _ in range(num_metrics_per_run)]
+    session.add_all(metrics)
+    session.commit()
+
+    print("🌳 seeding completed.")
+
+
+if __name__ == "__main__":
+    seed_all(
+        num_projects=3,
+        num_exps_per_project=20,
+        num_trials_per_exp=50,
+        num_runs_per_trial=100,
+        num_metrics_per_run=100,
+    )

alphatrion/utils/__init__.py

@@ -22,6 +22,7 @@ pytest = ">=8.4.2,<9.0.0"
 ruff = "^0.12.12"
 pytest-asyncio = ">=0.22.0,<1.0.0"
 pytest-timeout = ">=2.1.0,<3.0.0"
+faker = "^37.12.0"
 
 [build-system]
 requires = ["poetry-core>=2.0.0,<3.0.0"]