Add TreeManagerConfiguration.max_prompt_lottery_waiting parameter (#1889)

- add lang column to message_tree_state, create multi-column index on
`state`, `lang`
- stop generating initial prompt tasks for languages with more than
`max_prompt_lottery_waiting ` prompts in `prompt_lottery_waiting` state
- add `MAX_PROMPT_LOTTERY_WAITING` workflow variable

Author: andreaskoepf

.github/workflows/deploy-to-node.yaml

@@ -46,6 +46,7 @@ jobs:
       P_LONELY_CHILD_EXTENSION: ${{ vars.P_LONELY_CHILD_EXTENSION }}
       P_ACTIVATE_BACKLOG_TREE: ${{ vars.P_ACTIVATE_BACKLOG_TREE }}
       MIN_ACTIVE_RANKINGS_PER_LANG: ${{ vars.MIN_ACTIVE_RANKINGS_PER_LANG }}
+      MAX_PROMPT_LOTTERY_WAITING: ${{ vars.MAX_PROMPT_LOTTERY_WAITING }}
       GOAL_TREE_SIZE: ${{ vars.GOAL_TREE_SIZE }}
       MESSAGE_SIZE_LIMIT: ${{ vars.MESSAGE_SIZE_LIMIT }}
       SKIP_TOXICITY_CALCULATION: ${{ vars.SKIP_TOXICITY_CALCULATION }}

ansible/deploy-to-node.yaml

@@ -139,6 +139,9 @@
           TREE_MANAGER__MIN_ACTIVE_RANKINGS_PER_LANG:
             "{{ lookup('ansible.builtin.env', 'MIN_ACTIVE_RANKINGS_PER_LANG') |
             default('0', true) }}"
+          TREE_MANAGER__MAX_PROMPT_LOTTERY_WAITING:
+            "{{ lookup('ansible.builtin.env', 'MAX_PROMPT_LOTTERY_WAITING') |
+            default('250', true) }}"
           MESSAGE_SIZE_LIMIT:
             "{{ lookup('ansible.builtin.env', 'MESSAGE_SIZE_LIMIT') |
             default('2000', true) }}"

backend/alembic/versions/2023_02_26_0052-9db92d504f64_add_lang_to_message_tree_state.py

@@ -0,0 +1,35 @@
+"""add lang to message_tree_state
+
+Revision ID: 9db92d504f64
+Revises: 8cd0c34d0c3c
+Create Date: 2023-02-26 00:52:40.624843
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "9db92d504f64"
+down_revision = "8cd0c34d0c3c"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("message_tree_state", sa.Column("lang", sa.String(length=32), nullable=True))
+    op.execute(
+        "WITH msg AS (SELECT id, lang FROM message WHERE parent_id is NULL) UPDATE message_tree_state mts SET lang = msg.lang FROM msg WHERE mts.message_tree_id = msg.id"
+    )
+    op.alter_column("message_tree_state", "lang", nullable=False)
+    op.drop_index("ix_message_tree_state_state", table_name="message_tree_state")
+    op.create_index("ix_message_tree_state__lang__state", "message_tree_state", ["state", "lang"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("ix_message_tree_state__lang__state", table_name="message_tree_state")
+    op.create_index("ix_message_tree_state_state", "message_tree_state", ["state"], unique=False)
+    op.drop_column("message_tree_state", "lang")
+    # ### end Alembic commands ###

backend/oasst_backend/config.py

@@ -54,7 +54,7 @@ class TreeManagerConfiguration(BaseModel):
     p_full_labeling_review_prompt: float = 1.0
     """Probability of full text-labeling (instead of mandatory only) for initial prompts."""
 
-    p_full_labeling_review_reply_assistant: float = 0.5
+    p_full_labeling_review_reply_assistant: float = 1.0
     """Probability of full text-labeling (instead of mandatory only) for assistant replies."""
 
     p_full_labeling_review_reply_prompter: float = 0.25
@@ -145,6 +145,10 @@ class TreeManagerConfiguration(BaseModel):
     """Maximum number of pending tasks (neither canceled nor completed) by a single user within
     the time span defined by `recent_tasks_span_sec`."""
 
+    max_prompt_lottery_waiting: int = 250
+    """Maximum number of prompts in prompt_lottery_waiting state per language. If this value
+    is exceeded no new initial prompt tasks for that language are generated."""
+
 
 class Settings(BaseSettings):
     PROJECT_NAME: str = "open-assistant backend"

backend/oasst_backend/models/message_tree_state.py

@@ -5,7 +5,7 @@
 
 import sqlalchemy as sa
 import sqlalchemy.dialects.postgresql as pg
-from sqlmodel import Field, SQLModel
+from sqlmodel import Field, Index, SQLModel
 
 
 class State(str, Enum):
@@ -74,14 +74,16 @@ class State(str, Enum):
 
 class MessageTreeState(SQLModel, table=True):
     __tablename__ = "message_tree_state"
+    __table_args__ = (Index("ix_message_tree_state__lang__state", "state", "lang", unique=False),)
 
     message_tree_id: UUID = Field(
         sa_column=sa.Column(pg.UUID(as_uuid=True), sa.ForeignKey("message.id"), primary_key=True)
     )
     goal_tree_size: int = Field(nullable=False)
     max_depth: int = Field(nullable=False)
     max_children_count: int = Field(nullable=False)
-    state: str = Field(nullable=False, max_length=128, index=True)
+    state: str = Field(nullable=False, max_length=128)
     active: bool = Field(nullable=False, index=True)
     origin: str = Field(sa_column=sa.Column(sa.String(1024), nullable=True))
     won_prompt_lottery_date: Optional[datetime] = Field(sa_column=sa.Column(sa.DateTime(timezone=True), nullable=True))
+    lang: str = Field(sa_column=sa.Column(sa.String(32), nullable=False))

backend/oasst_backend/tree_manager.py

@@ -214,7 +214,7 @@ def _determine_task_availability_internal(
     ) -> dict[protocol_schema.TaskRequestType, int]:
         task_count_by_type: dict[protocol_schema.TaskRequestType, int] = {t: 0 for t in protocol_schema.TaskRequestType}
 
-        task_count_by_type[protocol_schema.TaskRequestType.initial_prompt] = num_missing_prompts
+        task_count_by_type[protocol_schema.TaskRequestType.initial_prompt] = max(0, num_missing_prompts)
 
         task_count_by_type[protocol_schema.TaskRequestType.prompter_reply] = len(
             list(filter(lambda x: x.parent_role == "assistant", extendible_parents))
@@ -256,13 +256,14 @@ def _prompt_lottery(self, lang: str, max_activate: int = 1) -> int:
 
         while True:
             stats = self.tree_counts_by_state_stats(lang=lang, only_active=True)
-
+            prompt_lottery_waiting = self.query_prompt_lottery_waiting(lang=lang)
+            remaining_lottery_entries = max(0, self.cfg.max_prompt_lottery_waiting - prompt_lottery_waiting)
             remaining_prompt_review = max(0, self.cfg.max_initial_prompt_review - stats.initial_prompt_review)
             num_missing_growing = max(0, self.cfg.max_active_trees - stats.growing)
             logger.info(f"_prompt_lottery {remaining_prompt_review=}, {num_missing_growing=}")
 
             if num_missing_growing == 0 or activated >= max_activate:
-                return num_missing_growing + remaining_prompt_review
+                return min(num_missing_growing + remaining_prompt_review, remaining_lottery_entries)
 
             @managed_tx_function(CommitMode.COMMIT)
             def activate_one(db: Session) -> int:
@@ -330,7 +331,7 @@ def activate_one(db: Session) -> int:
                 return True
 
             if not activate_one():
-                return num_missing_growing + remaining_prompt_review
+                return min(num_missing_growing + remaining_prompt_review, remaining_lottery_entries)
 
             activated += 1
 
@@ -714,8 +715,10 @@ async def handle_interaction(self, interaction: protocol_schema.AnyInteraction)
                 )
 
                 if not message.parent_id:
-                    logger.info(f"TreeManager: Inserting new tree state for initial prompt {message.id=}")
-                    self._insert_default_state(message.id)
+                    logger.info(
+                        f"TreeManager: Inserting new tree state for initial prompt {message.id=} [{message.lang}]"
+                    )
+                    self._insert_default_state(message.id, message.lang)
 
                 if not settings.DEBUG_SKIP_EMBEDDING_COMPUTATION:
                     try:
@@ -1261,10 +1264,10 @@ def query_tree_size(self, message_tree_id: UUID) -> ActiveTreeSizeRow:
 
         return ActiveTreeSizeRow.from_orm(qry.one())
 
-    def query_misssing_tree_states(self) -> list[UUID]:
+    def query_misssing_tree_states(self) -> list[Tuple[UUID, str]]:
         """Find all initial prompt messages that have no associated message tree state"""
         qry_missing_tree_states = (
-            self.db.query(Message.id)
+            self.db.query(Message.id, Message.lang)
             .outerjoin(MessageTreeState, Message.message_tree_id == MessageTreeState.message_tree_id)
             .filter(
                 Message.parent_id.is_(None),
@@ -1273,7 +1276,7 @@ def query_misssing_tree_states(self) -> list[UUID]:
             )
         )
 
-        return [m.id for m in qry_missing_tree_states.all()]
+        return [(m.id, m.lang) for m in qry_missing_tree_states.all()]
 
     _sql_find_tree_ranking_results = """
 -- get all ranking results of completed tasks for all parents with >= 2 children
@@ -1326,13 +1329,13 @@ def ensure_tree_states(self) -> None:
         """Add message tree state rows for all root nodes (initial prompt messages)."""
 
         missing_tree_ids = self.query_misssing_tree_states()
-        for id in missing_tree_ids:
+        for id, lang in missing_tree_ids:
             tree_size = self.db.query(func.count(Message.id)).filter(Message.message_tree_id == id).scalar()
             state = message_tree_state.State.INITIAL_PROMPT_REVIEW
             if tree_size > 1:
                 state = message_tree_state.State.GROWING
                 logger.info(f"Inserting missing message tree state for message: {id} ({tree_size=}, {state=:s})")
-            self._insert_default_state(id, state=state)
+            self._insert_default_state(id, lang=lang, state=state)
 
         halt_prompts_of_disabled_users(self.db)
 
@@ -1388,6 +1391,12 @@ def query_num_growing_trees(self, lang: str) -> int:
         )
         return query.scalar()
 
+    def query_prompt_lottery_waiting(self, lang: str) -> int:
+        query = self.db.query(func.count(MessageTreeState.message_tree_id)).filter(
+            MessageTreeState.state == message_tree_state.State.PROMPT_LOTTERY_WAITING, MessageTreeState.lang == lang
+        )
+        return query.scalar()
+
     def query_num_active_trees(
         self, lang: str, exclude_ranking: bool = True, exclude_prompt_review: bool = True
     ) -> int:
@@ -1451,6 +1460,7 @@ def _insert_tree_state(
         max_depth: int,
         max_children_count: int,
         active: bool,
+        lang: str,
         state: message_tree_state.State = message_tree_state.State.INITIAL_PROMPT_REVIEW,
     ) -> MessageTreeState:
         model = MessageTreeState(
@@ -1460,6 +1470,7 @@ def _insert_tree_state(
             max_children_count=max_children_count,
             state=state.value,
             active=active,
+            lang=lang,
         )
 
         self.db.add(model)
@@ -1469,6 +1480,7 @@ def _insert_tree_state(
     def _insert_default_state(
         self,
         root_message_id: UUID,
+        lang: str,
         state: message_tree_state.State = message_tree_state.State.INITIAL_PROMPT_REVIEW,
         *,
         goal_tree_size: int = None,
@@ -1483,8 +1495,9 @@ def _insert_default_state(
             goal_tree_size=goal_tree_size,
             max_depth=self.cfg.max_tree_depth,
             max_children_count=self.cfg.max_children_count,
-            state=state,
             active=True,
+            lang=lang,
+            state=state,
         )
 
     def tree_counts_by_state(self, lang: str = None, only_active: bool = False) -> dict[str, int]: