Design an Embedding and Classification API
Company: Scale AI
Role: Machine Learning Engineer
Category: ML System Design
Difficulty: easy
Interview Round: Onsite
## Design an Embedding and Classification API
Design a service that accepts text items, generates vector embeddings, and optionally classifies each item using a configured classifier built for a particular embedding model. Support both interactive single-item requests and larger batches. Begin by clarifying whether callers may supply existing embeddings; for the remainder, assume callers send raw text and the service owns preprocessing and embedding generation.
### Constraints & Assumptions
- Embedding and classifier versions must be explicit and compatible.
- A batch can contain valid and invalid items; define partial-failure semantics.
- Model execution may require scarce accelerators and benefits from batching.
- Tenants must not read one another's inputs, vectors, labels, or cached results.
- Do not assume request volume, text length, vector size, label count, latency, or freshness targets; identify how they affect the design.
### Part 1 — Define API and Version Contracts
Specify synchronous and batch APIs, input validation, outputs, idempotency, and version selection.
#### What This Part Should Cover
- Stable request, batch, and per-item identifiers.
- Text limits, model and classifier version, label taxonomy, scores, and per-item errors.
- A synchronous threshold and an asynchronous job contract for larger batches.
- Deterministic preprocessing metadata and explicit compatibility validation.
```hint Version the whole inference path
The same model name can produce different behavior if tokenization, preprocessing, classifier weights, or label definitions change.
```
### Part 2 — Build the Inference Path
Trace a request through preprocessing, embedding inference, and classification. Explain batching, caching, model loading, and how a classifier is tied to the embedding representation it expects.
#### What This Part Should Cover
- Stateless request handling separated from model-serving workers.
- Dynamic batching with a bounded wait and input-size-aware limits.
- A registry of immutable model artifacts and allowed compatibility pairs.
- Cache keys that include normalized input and all behavior-defining versions.
```hint A vector space is a contract
A classifier trained on one embedding version cannot silently consume vectors from another version merely because their dimensions match.
```
### Part 3 — Scale and Survive Failure
Design admission control, queues, worker pools, sharding, retry, and backpressure for interactive and batch traffic.
#### What This Part Should Cover
- Separate latency-sensitive and throughput-oriented capacity or scheduling.
- Tenant quotas, bounded queues, cancellation, and overload behavior.
- Idempotent batch retries and item-level checkpoints.
- Handling of worker crash, model-load failure, timeout, and partial batch failure.
```hint Do not retry a whole batch blindly
Persist item identity and status so completed inference is not repeated after one worker or item fails.
```
### Part 4 — Evaluate, Observe, and Roll Out Models
Explain offline and online quality checks, drift detection, privacy controls, observability, and safe version rollout.
#### What This Part Should Cover
- Representative embedding and classification evaluation sets with slice metrics.
- Shadow or canary comparison before activating a new compatible pair.
- Monitoring of latency, queue age, batch efficiency, errors, label distribution, and abstention.
- Retention, logging, and access choices for raw text, embeddings, and predictions.
```hint Separate service health from model quality
A fast successful response can still contain a degraded classification.
```
### What a Strong Answer Covers
- Precise API, version, compatibility, and partial-failure semantics.
- A concrete inference path with justified batching, caching, and resource isolation.
- Scale and retry behavior that does not duplicate completed batch work.
- Evaluation, rollout, privacy, drift, and operational signals for both embeddings and labels.
- Design choices tied to stated latency, throughput, and model assumptions.
### Follow-up Questions
1. How would the contract change if callers can submit their own embeddings?
2. When is caching an embedding unsafe or ineffective?
3. How would you prevent large batch traffic from starving interactive requests?
4. What evidence would block rollout of a new embedding-classifier pair even if latency improved?
Quick Answer: Design an API that generates text embeddings and optional compatible classifications for interactive and batch requests. Cover version contracts, preprocessing ownership, batching, accelerator use, partial failures, tenant isolation, caching, evaluation, observability, and safe model rollout.