Skip to content

Remove pytest from runtime model import#923

Merged
jimmyt857 merged 1 commit intoPhysical-Intelligence:mainfrom
kstonekuan:fix/remove-pytest-runtime-import
Apr 15, 2026
Merged

Remove pytest from runtime model import#923
jimmyt857 merged 1 commit intoPhysical-Intelligence:mainfrom
kstonekuan:fix/remove-pytest-runtime-import

Conversation

@kstonekuan
Copy link
Copy Markdown
Contributor

Summary

This removes an unintended runtime dependency on pytest from openpi.models_pytorch.gemma_pytorch.

The only usage was in a type annotation for past_key_values:

  • before: list[torch.FloatTensor] | pytest.Cache | None
  • after: list[torch.FloatTensor] | None

Why

pytest should not be required to import or serve runtime model code.

This showed up when using openpi in a production inference image that excludes dev dependencies.

In our serving environment, openpi.models_pytorch.gemma_pytorch was imported during model load, which caused the server to fail if pytest was not installed. The pytest reference was only used in a type annotation and was not part of runtime behavior.

Changes

  • removed import pytest from src/openpi/models_pytorch/gemma_pytorch.py
  • replaced the pytest.Cache annotation on past_key_values with a runtime-safe annotation

Validation

  • verified that the only runtime use of pytest in this module was the type annotation
  • verified the change is isolated to this import and annotation

Notes

I kept this patch intentionally narrow to avoid changing model behavior or cache semantics beyond removing the accidental test-only dependency.

@jimmyt857
Copy link
Copy Markdown
Contributor

Thanks!

@jimmyt857 jimmyt857 merged commit e4429ad into Physical-Intelligence:main Apr 15, 2026
2 checks passed
@kstonekuan kstonekuan deleted the fix/remove-pytest-runtime-import branch April 15, 2026 04:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants