[Bugfix][CI] Release the shared ColBERT engine before `test_colbert_hf_comparison` (#52608)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Aug 17, 2026 at 19:50 UTC
402547d7f02bdbfc5dce5d27dc21f50dd4d627b6
1 file changed
+71
-82
tests/models/language/pooling/test_colbert.py
+71
-82
@@ -186,33 +186,33 @@ def _assert_embeddings_close(vllm_outputs, hf_embeddings):
186
)
187
188
189
-@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="module")
189
+@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="class")
190
def colbert_spec(request):
191
"""Return the model spec dict for the current parametrization."""
192
return COLBERT_MODELS[request.param]
193
194
195
-@pytest.fixture(scope="module")
195
+@pytest.fixture(scope="class")
196
def colbert_model_name(colbert_spec):
197
return colbert_spec["model"]
198
199
200
-@pytest.fixture(scope="module")
200
+@pytest.fixture(scope="class")
201
def colbert_dim(colbert_spec):
202
return colbert_spec["colbert_dim"]
203
204
205
-@pytest.fixture(scope="module")
205
+@pytest.fixture(scope="class")
206
def colbert_max_model_len(colbert_spec):
207
return colbert_spec["max_model_len"]
208
209
210
-@pytest.fixture(scope="module")
210
+@pytest.fixture(scope="class")
211
def colbert_extra_kwargs(colbert_spec):
212
return colbert_spec["extra_kwargs"]
213
214
215
-@pytest.fixture(scope="module")
215
+@pytest.fixture(scope="class")
216
def colbert_model(
217
vllm_runner,
218
colbert_model_name,
@@ -230,103 +230,92 @@ def colbert_model(
230
yield vllm_model
231
232
233
-def test_colbert_token_embed(
234
- colbert_model,
235
- colbert_dim,
236
-):
237
- """Test that ColBERT model produces token embeddings."""
238
- outputs = colbert_model.token_embed([TEXTS_1[0]])
239
-
240
- assert len(outputs) == 1
241
- emb = torch.as_tensor(outputs[0])
242
- assert emb.dim() == 2
243
- assert emb.shape[1] == colbert_dim
244
- assert emb.shape[0] > 1
245
-
246
-
247
-def test_colbert_late_interaction_1_to_1(
248
- colbert_model,
249
-):
250
- """Test ColBERT late interaction scoring with 1:1 query-document pair."""
251
- q_outputs = colbert_model.token_embed([TEXTS_1[0]])
252
- d_outputs = colbert_model.token_embed([TEXTS_2[0]])
233
+class TestColbertSharedEngine:
234
+ """Tests sharing one engine per model.
235
254
- q_emb = torch.as_tensor(q_outputs[0])
255
- d_emb = torch.as_tensor(d_outputs[0])
236
+ Class-scoped so the engine is released before `test_colbert_hf_comparison`,
237
+ which needs a runner of its own and cannot start while this one holds VRAM.
238
+ """
239
257
- manual_score = compute_maxsim_score(q_emb, d_emb).item()
240
+ def test_colbert_token_embed(self, colbert_model, colbert_dim):
241
+ """Test that ColBERT model produces token embeddings."""
242
+ outputs = colbert_model.token_embed([TEXTS_1[0]])
243
259
- vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0])
244
+ assert len(outputs) == 1
245
+ emb = torch.as_tensor(outputs[0])
246
+ assert emb.dim() == 2
247
+ assert emb.shape[1] == colbert_dim
248
+ assert emb.shape[0] > 1
249
261
- assert len(vllm_scores) == 1
262
- assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
250
+ def test_colbert_late_interaction_1_to_1(self, colbert_model):
251
+ """Test ColBERT late interaction scoring with 1:1 query-document pair."""
252
+ q_outputs = colbert_model.token_embed([TEXTS_1[0]])
253
+ d_outputs = colbert_model.token_embed([TEXTS_2[0]])
254
255
+ q_emb = torch.as_tensor(q_outputs[0])
256
+ d_emb = torch.as_tensor(d_outputs[0])
257
265
-def test_colbert_late_interaction_1_to_N(
266
- colbert_model,
267
-):
268
- """Test ColBERT late interaction scoring with 1:N query-documents."""
269
- q_outputs = colbert_model.token_embed([TEXTS_1[0]])
270
- d_outputs = colbert_model.token_embed(TEXTS_2)
258
+ manual_score = compute_maxsim_score(q_emb, d_emb).item()
259
272
- q_emb = torch.as_tensor(q_outputs[0])
260
+ vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2[0])
261
274
- manual_scores = []
275
- for d_out in d_outputs:
276
- d_emb = torch.as_tensor(d_out)
277
- manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
262
+ assert len(vllm_scores) == 1
263
+ assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
264
279
- vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2)
265
+ def test_colbert_late_interaction_1_to_N(self, colbert_model):
266
+ """Test ColBERT late interaction scoring with 1:N query-documents."""
267
+ q_outputs = colbert_model.token_embed([TEXTS_1[0]])
268
+ d_outputs = colbert_model.token_embed(TEXTS_2)
269
281
- assert len(vllm_scores) == 2
282
- for i in range(2):
283
- assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
270
+ q_emb = torch.as_tensor(q_outputs[0])
271
272
+ manual_scores = []
273
+ for d_out in d_outputs:
274
+ d_emb = torch.as_tensor(d_out)
275
+ manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
276
286
-def test_colbert_late_interaction_N_to_N(
287
- colbert_model,
288
-):
289
- """Test ColBERT late interaction scoring with N:N query-documents."""
290
- q_outputs = colbert_model.token_embed(TEXTS_1)
291
- d_outputs = colbert_model.token_embed(TEXTS_2)
277
+ vllm_scores = colbert_model.score(TEXTS_1[0], TEXTS_2)
278
293
- manual_scores = []
294
- for q_out, d_out in zip(q_outputs, d_outputs):
295
- q_emb = torch.as_tensor(q_out)
296
- d_emb = torch.as_tensor(d_out)
297
- manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
279
+ assert len(vllm_scores) == 2
280
+ for i in range(2):
281
+ assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
282
299
- vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2)
283
+ def test_colbert_late_interaction_N_to_N(self, colbert_model):
284
+ """Test ColBERT late interaction scoring with N:N query-documents."""
285
+ q_outputs = colbert_model.token_embed(TEXTS_1)
286
+ d_outputs = colbert_model.token_embed(TEXTS_2)
287
301
- assert len(vllm_scores) == 2
302
- for i in range(2):
303
- assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
288
+ manual_scores = []
289
+ for q_out, d_out in zip(q_outputs, d_outputs):
290
+ q_emb = torch.as_tensor(q_out)
291
+ d_emb = torch.as_tensor(d_out)
292
+ manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
293
294
+ vllm_scores = colbert_model.score(TEXTS_1, TEXTS_2)
295
306
-def test_colbert_relevance_ordering(
307
- colbert_model,
308
-):
309
- """Test that ColBERT scores relevant documents higher than irrelevant."""
310
- query = "What is machine learning?"
311
- documents = [
312
- "Machine learning is a subset of artificial intelligence.",
313
- "Python is a programming language.",
314
- "Deep learning uses neural networks.",
315
- ]
296
+ assert len(vllm_scores) == 2
297
+ for i in range(2):
298
+ assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
299
317
- scores = colbert_model.score(query, documents)
300
+ def test_colbert_relevance_ordering(self, colbert_model):
301
+ """Test that ColBERT scores relevant documents higher than irrelevant."""
302
+ query = "What is machine learning?"
303
+ documents = [
304
+ "Machine learning is a subset of artificial intelligence.",
305
+ "Python is a programming language.",
306
+ "Deep learning uses neural networks.",
307
+ ]
308
319
- assert len(scores) == 3
320
- assert scores[0] > scores[1], "ML doc should score higher than Python doc"
321
- assert scores[2] > scores[1], "DL doc should score higher than Python doc"
309
+ scores = colbert_model.score(query, documents)
310
311
+ assert len(scores) == 3
312
+ assert scores[0] > scores[1], "ML doc should score higher than Python doc"
313
+ assert scores[2] > scores[1], "DL doc should score higher than Python doc"
314
324
-def test_colbert_embed_not_supported(
325
- colbert_model,
326
-):
327
- """Test that ColBERT model does not support the embed task."""
328
- with pytest.raises(ValueError, match="Embedding API is not supported"):
329
- colbert_model.embed([TEXTS_1[0]])
315
+ def test_colbert_embed_not_supported(self, colbert_model):
316
+ """Test that ColBERT model does not support the embed task."""
317
+ with pytest.raises(ValueError, match="Embedding API is not supported"):
318
+ colbert_model.embed([TEXTS_1[0]])
319
320
321
@pytest.mark.parametrize(