Create and clean up kafka topics manually to reduce memory burden for Adala server

docker-compose.yml

@@ -20,6 +20,7 @@ services:
  - ALLOW_PLAINTEXT_LISTENER=yes
  - KAFKA_CFG_NODE_ID=1
  - KAFKA_KRAFT_CLUSTER_ID=MkU3OEVBNTcwNTJENDM2Qk
+ - KAFKA_CFG_AUTO_CREATE_TOPICS_ENABLE=false
  app:
  build:
  context: .
@@ -30,8 +31,9 @@ services:
  redis:
  condition: service_healthy
  environment:
- - KAFKA_BOOTSTRAP_SERVERS=kafka:9093
  - REDIS_URL=redis://redis:6379/0
+ - KAFKA_BOOTSTRAP_SERVERS=kafka:9093 # TODO pull from .env
+ - KAFKA_RETENTION_MS=180000 # TODO pull from .env
  command:
  ["poetry", "run", "uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
  worker:

pyproject.toml

@@ -36,6 +36,7 @@ celery = {version = "^5.3.6", extras = ["redis"]}
 uvicorn = "*"
 pydantic-settings = "^2.2.1"
 label-studio-sdk = "^0.0.32"
+kafka-python = "^2.0.2"
 
 [tool.poetry.dev-dependencies]
 pytest = "^7.4.3"

server/.env.example

@@ -1 +1,4 @@
 KAFKA_BOOTSTRAP_SERVERS="localhost:9093"
+
+# this value is only for local dev. In our deployments, it is not set here, but in another place: https://github.com/HumanSignal/infra/pull/67
+KAFKA_RETENTION_MS=180000 # 30 minutes

server/app.py

@@ -25,7 +25,7 @@
  process_streaming_output,
  streaming_parent_task,
 )
-from utils import get_input_topic, Settings
+from utils import get_input_topic_name, Settings
 from server.handlers.result_handlers import ResultHandler
 
 
@@ -230,7 +230,7 @@ async def submit_batch(batch: BatchData):
  Response: Generic response indicating status of request
  """
 
- topic = get_input_topic(batch.job_id)
+ topic = get_input_topic_name(batch.job_id)
  producer = AIOKafkaProducer(
  bootstrap_servers=settings.kafka_bootstrap_servers,
  value_serializer=lambda v: json.dumps(v).encode("utf-8"),

server/tasks/process_file.py

@@ -11,7 +11,13 @@
 from aiokafka.errors import UnknownTopicOrPartitionError
 from celery import Celery, states
 from celery.exceptions import Ignore
-from server.utils import get_input_topic, get_output_topic, Settings
+from server.utils import (
+ get_input_topic_name,
+ get_output_topic_name,
+ ensure_topic,
+ delete_topic,
+ Settings,
+)
 from server.handlers.result_handlers import ResultHandler
 
 
@@ -54,21 +60,29 @@ def streaming_parent_task(
  # Parent job ID is used for input/output topic names
  parent_job_id = self.request.id
 
- # Override kafka_bootstrap_servers with value from settings
+ # create kafka topics
+ input_topic_name = get_input_topic_name(parent_job_id)
+ ensure_topic(input_topic_name)
+ output_topic_name = get_output_topic_name(parent_job_id)
+ ensure_topic(output_topic_name)
+
+ # Override default agent kafka settings
  settings = Settings()
  agent.environment.kafka_bootstrap_servers = settings.kafka_bootstrap_servers
+ agent.environment.kafka_input_topic = input_topic_name
+ agent.environment.kafka_output_topic = output_topic_name
 
  inference_task = process_file_streaming
  logger.info(f"Submitting task {inference_task.name} with agent {agent}")
- input_result = inference_task.delay(agent=agent, parent_job_id=parent_job_id)
+ input_result = inference_task.delay(agent=agent)
  input_job_id = input_result.id
  logger.info(f"Task {inference_task.name} submitted with job_id {input_job_id}")
 
  result_handler_task = process_streaming_output
  logger.info(f"Submitting task {result_handler_task.name}")
  output_result = result_handler_task.delay(
  input_job_id=input_job_id,
- parent_job_id=parent_job_id,
+ output_topic_name=output_topic_name,
  result_handler=result_handler,
  batch_size=batch_size,
  )
@@ -95,6 +109,10 @@ def streaming_parent_task(
  ):
  time.sleep(1)
 
+ # clean up kafka topics
+ delete_topic(input_topic_name)
+ delete_topic(output_topic_name)
+
  logger.info("Both input and output jobs complete")
 
  # Update parent task status to SUCCESS and pass metadata again
@@ -109,45 +127,40 @@ def streaming_parent_task(
  raise Ignore()
 
 
-@app.task(
- name="process_file_streaming", track_started=True, bind=True, serializer="pickle"
-)
-def process_file_streaming(self, agent: Agent, parent_job_id: str):
- # Set input and output topics using parent job ID
- agent.environment.kafka_input_topic = get_input_topic(parent_job_id)
- agent.environment.kafka_output_topic = get_output_topic(parent_job_id)
+@app.task(name="process_file_streaming", track_started=True, serializer="pickle")
+def process_file_streaming(agent: Agent):
+ # agent's kafka_bootstrap servers and kafka topics should be set in parent task
 
  # Run the agent
  asyncio.run(agent.arun())
 
 
 async def async_process_streaming_output(
  input_job_id: str,
- parent_job_id: str,
+ output_topic_name,
  result_handler: ResultHandler,
  batch_size: int,
 ):
- logger.info(f"Polling for results {parent_job_id=}")
+ logger.info(f"Polling for results {output_topic_name=}")
 
- topic = get_output_topic(parent_job_id)
  settings = Settings()
 
  # Retry to workaround race condition of topic creation
  retries = 5
  while retries > 0:
  try:
  consumer = AIOKafkaConsumer(
- topic,
+ output_topic_name,
  bootstrap_servers=settings.kafka_bootstrap_servers,
  value_deserializer=lambda v: json.loads(v.decode("utf-8")),
  auto_offset_reset="earliest",
  )
  await consumer.start()
- logger.info(f"consumer started {parent_job_id=}")
+ logger.info(f"consumer started {output_topic_name=}")
  break
  except UnknownTopicOrPartitionError as e:
  logger.error(msg=e)
- logger.info(f"Retrying to create consumer with topic {topic}")
+ logger.info(f"Retrying to create consumer with topic {output_topic_name}")
 
  await consumer.stop()
  retries -= 1
@@ -183,20 +196,17 @@ async def async_process_streaming_output(
  await consumer.stop()
 
 
-@app.task(
- name="process_streaming_output", track_started=True, bind=True, serializer="pickle"
-)
+@app.task(name="process_streaming_output", track_started=True, serializer="pickle")
 def process_streaming_output(
- self,
  input_job_id: str,
- parent_job_id: str,
+ output_topic_name: str,
  result_handler: ResultHandler,
  batch_size: int,
 ):
  try:
  asyncio.run(
  async_process_streaming_output(
- input_job_id, parent_job_id, result_handler, batch_size
+ input_job_id, output_topic_name, result_handler, batch_size
  )
  )
  except Exception as e:

server/utils.py

@@ -1,6 +1,8 @@
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from typing import List, Union
 from pathlib import Path
+from kafka.admin import KafkaAdminClient, NewTopic
+from kafka.errors import TopicAlreadyExistsError
 
 
 class Settings(BaseSettings):
@@ -10,16 +12,55 @@ class Settings(BaseSettings):
  """
 
  kafka_bootstrap_servers: Union[str, List[str]]
+ kafka_retention_ms: int
 
  model_config = SettingsConfigDict(
  # have to use an absolute path here so celery workers can find it
  env_file=(Path(__file__).parent / ".env"),
  )
 
 
-def get_input_topic(job_id: str):
- return f"adala-input-{job_id}"
+def get_input_topic_name(job_id: str):
+ topic_name = f"adala-input-{job_id}"
 
+ return topic_name
 
-def get_output_topic(job_id: str):
- return f"adala-output-{job_id}"
+
+def get_output_topic_name(job_id: str):
+ topic_name = f"adala-output-{job_id}"
+
+ return topic_name
+
+
+def ensure_topic(topic_name: str):
+ settings = Settings()
+ bootstrap_servers = settings.kafka_bootstrap_servers
+ retention_ms = settings.kafka_retention_ms
+
+ admin_client = KafkaAdminClient(
+ bootstrap_servers=bootstrap_servers, client_id="topic_creator"
+ )
+
+ topic = NewTopic(
+ name=topic_name,
+ num_partitions=1,
+ replication_factor=1,
+ topic_configs={"retention.ms": str(retention_ms)},
+ )
+
+ try:
+ admin_client.create_topics(new_topics=[topic])
+ except TopicAlreadyExistsError:
+ # we shouldn't hit this case when KAFKA_CFG_AUTO_CREATE_TOPICS=false unless there is a legitimate name collision, so should raise here after testing
+ pass
+
+
+def delete_topic(topic_name: str):
+ settings = Settings()
+ bootstrap_servers = settings.kafka_bootstrap_servers
+
+ admin_client = KafkaAdminClient(
+ bootstrap_servers=bootstrap_servers, client_id="topic_deleter"
+ )
+
+ admin_client.delete_topics(topics=[topic_name])