Simplify FFT request processor boundary

Author: ShubyM

src/server/gateway.py

@@ -15,20 +15,7 @@
 from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
-from pydantic import BaseModel
 from store import get_store
-from training_request_types import (
-  CreateModelFromStatePayload,
-  CreateModelPayload,
-  ForwardBackwardPayload,
-  LoadWeightsPayload,
-  OptimStepPayload,
-  SamplePayload,
-  SaveStatePayload,
-  SaveWeightsForSamplerPayload,
-  TrainingCommand,
-  TrainingOp,
-)
 from worker_launch_processor import (
   FFTWorkerManager,
   WorkerLaunchProcessor,
@@ -122,41 +109,40 @@ def is_sampler_weights_ref(model_id: str | None) -> bool:
   return len(parts) >= 3 and parts[1] == "sampler_weights"
 
 
-def make_training_command(
-  op: TrainingOp,
+def make_training_request(
+  op: str,
   model_id: str | None,
-  payload: BaseModel,
+  payload: dict,
   request_id: str | None = None,
-) -> TrainingCommand:
-  return TrainingCommand(
-    request_id=request_id or str(uuid.uuid4()),
-    op=op,
-    model_id=model_id,
-    payload=payload.model_dump(exclude_none=True),
-  )
+) -> dict:
+  request = {
+    "request_id": request_id or str(uuid.uuid4()),
+    "op": op,
+    "payload": payload,
+  }
+  if model_id is not None:
+    request["model_id"] = model_id
+  return request
 
 
-async def prepare_enqueue(command: TrainingCommand) -> tuple[str, dict]:
+async def enqueue(request: dict) -> str:
+  """Create a pending future, inject trace context, push to store. Returns req_id."""
+  request_id = request["request_id"]
   carrier: dict = {}
   propagate.inject(carrier)
-  command = command.model_copy(update={"trace_context": carrier})
-  payload = command.model_dump(exclude_none=True)
-  await store.set_future(command.request_id, {"status": "pending"})
-  return command.request_id, payload
-
-
-async def enqueue(command: TrainingCommand) -> str:
-  """Create a pending future, inject trace context, push to store. Returns req_id."""
-  req_id, payload = await prepare_enqueue(command)
-  await store.put_request(payload)
-  return req_id
+  await store.set_future(request_id, {"status": "pending"})
+  await store.put_request({**request, "trace_context": carrier})
+  return request_id
 
 
-async def enqueue_worker_launch(command: TrainingCommand) -> str:
+async def enqueue_worker_launch(request: dict) -> str:
   """Create a pending future and push a create-model request to the worker launch queue."""
-  req_id, payload = await prepare_enqueue(command)
-  await store.put_worker_launch_request(payload)
-  return req_id
+  request_id = request["request_id"]
+  carrier: dict = {}
+  propagate.inject(carrier)
+  await store.set_future(request_id, {"status": "pending"})
+  await store.put_worker_launch_request({**request, "trace_context": carrier})
+  return request_id
 
 
 async def preflight_vllm() -> None:
@@ -236,10 +222,10 @@ async def lifespan(_: FastAPI):
     if not is_fft_enabled():
       import training_requests_processor
 
-      worker = training_requests_processor.create_training_worker()
+      trainer = training_requests_processor.LoraTrainingWorker()
       if base_model:
-        await asyncio.to_thread(worker.load_base_model, base_model)
-      task = asyncio.create_task(training_requests_processor.run_training_requests_processor(worker))
+        await asyncio.to_thread(trainer.load_base_model, base_model)
+      task = asyncio.create_task(training_requests_processor.run_training_requests_processor(trainer))
   try:
     yield
   finally:
@@ -298,14 +284,14 @@ async def create_model(req: dict):
   if not base_model:
     return JSONResponse(status_code=400, content={"error": "base_model is required"})
   model_id = str(uuid.uuid4())
-  command = make_training_command(
+  command = make_training_request(
     "create_model",
     model_id,
-    CreateModelPayload(
-      base_model=base_model,
-      lora_config=req.get("lora_config") or {},
-      full_config=req.get("full_config") or {},
-    ),
+    {
+      "base_model": base_model,
+      "lora_config": req.get("lora_config") or {},
+      "full_config": req.get("full_config") or {},
+    },
     request_id=model_id,
   )
   req_id = await enqueue_worker_launch(command) if is_fft_enabled() else await enqueue(command)
@@ -321,13 +307,13 @@ async def create_model_from_state(req: dict):
   # Resolve relative names under TMP_DIR/checkpoints, leave absolute paths alone.
   resolved_path = state_path if os.path.isabs(state_path) else os.path.join(TMP_DIR, "checkpoints", state_path)
   model_id = str(uuid.uuid4())
-  command = make_training_command(
+  command = make_training_request(
     "create_model_from_state",
     model_id,
-    CreateModelFromStatePayload(
-      state_path=resolved_path,
-      restore_optimizer=bool(req.get("restore_optimizer", False)),
-    ),
+    {
+      "state_path": resolved_path,
+      "restore_optimizer": bool(req.get("restore_optimizer", False)),
+    },
     request_id=model_id,
   )
   req_id = await enqueue_worker_launch(command) if is_fft_enabled() else await enqueue(command)
@@ -376,14 +362,14 @@ async def forward_backward(req: dict):
   """TrainingClient.forward_backward_async()"""
   fwd_input = req.get("forward_backward_input", {})
   req_id = await enqueue(
-    make_training_command(
+    make_training_request(
       "forward_backward",
       req.get("model_id"),
-      ForwardBackwardPayload(
-        data=fwd_input.get("data", []),
-        loss_fn=fwd_input.get("loss_fn", "cross_entropy"),
-        loss_config=fwd_input.get("loss_fn_config", {}),
-      ),
+      {
+        "data": fwd_input.get("data", []),
+        "loss_fn": fwd_input.get("loss_fn", "cross_entropy"),
+        "loss_config": fwd_input.get("loss_fn_config", {}),
+      },
     )
   )
   return {"request_id": req_id}
@@ -393,10 +379,10 @@ async def forward_backward(req: dict):
 async def optim_step(req: dict):
   """TrainingClient.optim_step_async()"""
   req_id = await enqueue(
-    make_training_command(
+    make_training_request(
       "optim_step",
       req.get("model_id"),
-      OptimStepPayload(adam_params=req.get("adam_params", {})),
+      {"adam_params": req.get("adam_params", {})},
     )
   )
   return {"request_id": req_id}
@@ -419,14 +405,14 @@ async def save_weights_for_sampler(req: dict):
 
   session_id = sampler_session_id(model_id, seq_id)
   req_id = await enqueue(
-    make_training_command(
+    make_training_request(
       "save_weights_for_sampler",
       model_id,
-      SaveWeightsForSamplerPayload(
-        alias=alias,
-        path=sampler_weights_path(model_id, alias) if alias else None,
-        sampling_session_id=session_id,
-      ),
+      {
+        "alias": alias,
+        "path": sampler_weights_path(model_id, alias) if alias else None,
+        "sampling_session_id": session_id,
+      },
     )
   )
   return {"request_id": req_id}
@@ -451,14 +437,14 @@ async def save_weights(req: dict):
 
   req_id = str(uuid.uuid4())
   await enqueue(
-    make_training_command(
+    make_training_request(
       "save_state",
       model_id,
-      SaveStatePayload(
-        state_path=state_path,
-        include_optimizer=bool(req.get("include_optimizer", False)),
-        kind="weights",
-      ),
+      {
+        "state_path": state_path,
+        "include_optimizer": bool(req.get("include_optimizer", False)),
+        "kind": "weights",
+      },
       request_id=req_id,
     )
   )
@@ -477,13 +463,13 @@ async def load_weights(req: dict):
 
   resolved_path = checkpoint_state_path(model_id, state_path)
   req_id = await enqueue(
-    make_training_command(
+    make_training_request(
       "load_weights",
       model_id,
-      LoadWeightsPayload(
-        state_path=resolved_path,
-        restore_optimizer=bool(req.get("optimizer", False)),
-      ),
+      {
+        "state_path": resolved_path,
+        "restore_optimizer": bool(req.get("optimizer", False)),
+      },
     )
   )
   return {"request_id": req_id}
@@ -528,16 +514,16 @@ async def asample(req: dict):
 
   if get_sampler_backend() == "torch":
     req_id = await enqueue(
-      make_training_command(
+      make_training_request(
         "sample",
         base_model_id or model_id,
-        SamplePayload(
-          prompt_tokens=prompt,
-          max_tokens=max_tokens,
-          temperature=temperature,
-          num_samples=num_samples,
-          prompt_logprobs=bool(include_prompt_logprobs),
-        ),
+        {
+          "prompt_tokens": prompt,
+          "max_tokens": max_tokens,
+          "temperature": temperature,
+          "num_samples": num_samples,
+          "prompt_logprobs": bool(include_prompt_logprobs),
+        },
       )
     )
     return {"request_id": req_id}

src/server/training/trainer_worker.py

@@ -5,7 +5,7 @@
 from typing import Any
 
 import torch
-from pydantic import BaseModel, field_validator, model_validator
+from pydantic import BaseModel
 from transformers import PreTrainedModel, PreTrainedTokenizerBase
 
 from training import losses
@@ -14,32 +14,11 @@
 class TensorData(BaseModel):
   data: list[int] | list[float]
 
-  @model_validator(mode="before")
-  @classmethod
-  def accept_raw_sequence(cls, value):
-    if isinstance(value, list):
-      return {"data": value}
-    return value
-
 
 class Datum(BaseModel):
   loss_fn_inputs: dict[str, TensorData]
   model_input: list[int]
 
-  @field_validator("model_input", mode="before")
-  @classmethod
-  def flatten_model_input(cls, value):
-    if not isinstance(value, dict):
-      return value
-
-    tokens: list[int] = []
-    for chunk in value.get("chunks", []):
-      if isinstance(chunk, dict):
-        tokens.extend(chunk.get("tokens", []))
-      else:
-        tokens.extend(getattr(chunk, "tokens", []))
-    return tokens
-
 
 class BaseTrainerWorker:
   def __init__(self):