Skip to content

Commit 94ccd02

Browse files
authored
Fix forwarding ARN sequence lengths from QairtPipelinePass to QairtEncapsulationPass (microsoft#2526)
QairtEncapsulation uses sequence_lengths from model_attributes to compute max_length in genai_config.json. QairtGenAIBuilder populates this via its own config, but QairtPipelinePass had no equivalent path. Extract the arn list from the genai_builder.transform_options in the YAML recipe and pass it through model_attributes so QairtEncapsulation can pick it up. Previously caused an issue in execution where context length would be exceeded and some kind of crash would occur ## Describe your changes ## Checklist before requesting a review - [x] Add unit tests for this change. - [x] Make sure all tests can pass. - [x] Update documents if necessary. - [x] Lint and apply fixes to your code by running `lintrunner -a` - [x] Is this a user-facing change? If yes, give a description of this change to be included in the release notes. ## (Optional) Issue link
1 parent 449c7fc commit 94ccd02

2 files changed

Lines changed: 82 additions & 1 deletion

File tree

olive/passes/qairt/pipeline.py

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,4 +154,12 @@ def _run_for_config(
154154
if src.exists() and not dst.exists():
155155
shutil.copy2(src, dst)
156156

157-
return QairtModelHandler(model_path=output_model_path)
157+
# Forward AR decode lengths to QairtEncapsulation so it can compute max_length correctly.
158+
# arn in the pipeline recipe is the equivalent of sequence_lengths in QairtGenAIBuilder.
159+
arn = recipe_data.get("stages", {}).get("genai_builder", {}).get("transform_options", {}).get("arn")
160+
if arn is not None and (not isinstance(arn, list) or not arn or not all(isinstance(v, int) for v in arn)):
161+
raise ValueError(
162+
f"stages.genai_builder.transform_options.arn must be a non-empty list of integers, got: {arn!r}"
163+
)
164+
model_attrs = {"sequence_lengths": arn} if arn else {}
165+
return QairtModelHandler(model_path=output_model_path, model_attributes=model_attrs)

test/passes/qairt/test_pipeline_pass.py

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,79 @@ def test_pipeline_pass_missing_recipe_file(tmp_path, mock_hf_model, mock_pipelin
228228
pipeline_pass.run(mock_hf_model, str(output_path))
229229

230230

231+
def test_pipeline_pass_forwards_arn_as_sequence_lengths(tmp_path, mock_hf_model, recipe_file, mock_pipeline_modules):
232+
"""ARN from genai_builder.transform_options is forwarded as sequence_lengths in model_attributes."""
233+
output_path = tmp_path / "output"
234+
235+
mock_pipeline_modules["Recipe"].from_file.return_value = {
236+
"stages": {
237+
"genai_builder": {
238+
"transform_options": {
239+
"arn": [128, 1024, 4096],
240+
}
241+
}
242+
}
243+
}
244+
245+
pipeline_pass = create_pass_from_dict(
246+
QairtPipelinePass,
247+
{"recipe": str(recipe_file)},
248+
disable_search=True,
249+
)
250+
251+
result = pipeline_pass.run(mock_hf_model, str(output_path))
252+
253+
assert isinstance(result, QairtModelHandler)
254+
assert result.model_attributes == {"sequence_lengths": [128, 1024, 4096]}
255+
256+
257+
@pytest.mark.parametrize(
258+
"bad_arn",
259+
[
260+
"128,1024", # string instead of list
261+
1024, # bare int
262+
[], # empty list
263+
[128, "1024"], # list with non-int element
264+
],
265+
)
266+
def test_pipeline_pass_invalid_arn_raises(tmp_path, mock_hf_model, recipe_file, mock_pipeline_modules, bad_arn):
267+
"""A malformed arn value raises ValueError with a clear message before propagating downstream."""
268+
output_path = tmp_path / "output"
269+
270+
mock_pipeline_modules["Recipe"].from_file.return_value = {
271+
"stages": {"genai_builder": {"transform_options": {"arn": bad_arn}}}
272+
}
273+
274+
pipeline_pass = create_pass_from_dict(
275+
QairtPipelinePass,
276+
{"recipe": str(recipe_file)},
277+
disable_search=True,
278+
)
279+
280+
with pytest.raises(ValueError, match="must be a non-empty list of integers"):
281+
pipeline_pass.run(mock_hf_model, str(output_path))
282+
283+
284+
def test_pipeline_pass_no_arn_yields_empty_model_attributes(
285+
tmp_path, mock_hf_model, recipe_file, mock_pipeline_modules
286+
):
287+
"""When genai_builder.transform_options has no arn, sequence_lengths is not set in model_attributes."""
288+
output_path = tmp_path / "output"
289+
290+
mock_pipeline_modules["Recipe"].from_file.return_value = {"stages": {"genai_builder": {"transform_options": {}}}}
291+
292+
pipeline_pass = create_pass_from_dict(
293+
QairtPipelinePass,
294+
{"recipe": str(recipe_file)},
295+
disable_search=True,
296+
)
297+
298+
result = pipeline_pass.run(mock_hf_model, str(output_path))
299+
300+
assert isinstance(result, QairtModelHandler)
301+
assert "sequence_lengths" not in (result.model_attributes or {})
302+
303+
231304
def test_pipeline_pass_import_error(tmp_path, mock_hf_model, recipe_file):
232305
"""Test that ImportError is raised if qairt cannot be imported."""
233306

0 commit comments

Comments
 (0)