RequirementsManager rollback does not preserve non-PEP440 install sources (-e / @), causing environment drift after tests

[ashokkumarkannan1]

Description

• RequirementsManager is designed to temporarily install per-model dependencies and restore the environment after test execution.

• However, it currently tracks and restores only packages parsed from pip freeze entries of the form: name==version

• Editable installs (-e ...) and direct reference installs (name @ ...) are intentionally excluded from tracking. During rollback, changed packages are restored using: pip install name==version

• This leads to incomplete environment restoration when packages were originally installed via -e or @.

Current Behavior:

a) During __enter__:
- Snapshot pip freeze
- Install model-specific requirements
- Snapshot pip freeze again
- Compute diffs

b) _parse_freeze():
- Records only name==version entries.
- Skips -e and name @ ... entries (except special-cased handling)

c) During __exit__:
- Changed packages are restored using:
pip install <name>==<version>

Problem: Missing Tracking of -e / @ Installs

• If a package in the original environment was installed via:
  a) Editable install (-e path)
  b) Direct reference (name @ git+https://...)

• It is not properly tracked in _before_freeze / _after_freeze.

As a result:

• The package may not be restored to its original source.
• The environment remains mutated after test execution.

Incorrect Rollback Strategy

• Even if we start tracking -e / @ entries, the current rollback logic still restores using: pip install name==version
• This is incorrect for packages originally installed via: Git URLs, Editable installs, Direct references
• Because their original installation was not version-pin based, restoring via == changes their provenance.

Why This Is a Bug

• The context manager contract states that it: Temporarily installs requirements and rolls back changes after test execution.

• Today, rollback is incomplete for non-== package specs and can leave persistent environment mutations across:

Steps to reproduce the issue:

Case 1 — GPT-2 (easydel introduced)

pip show easydel
# WARNING: Package(s) not found
#Run test:
pytest tests/runner/test_models.py::test_all_models_jax[gpt2/causal_lm/jax-Base-single_device-inference] -svv


# After test:
pip show easydel
# Expected Output: Name: easydel, Version: 0.1.5.dev22, Location: .../site-packages
# easydel remains installed after the test — environment not rolled back.

Case 2 — pi_0 (transformers version drift)

# Initial environment: pip show transformers, Version: 4.57.1 (current)
# Run test:
pytest -s tests/runner/test_models.py::test_all_models_torch[pi_0/pytorch-lerobot_pi0_libero_base-single_device-inference]


#After test:
pip show transformers
# Expected Output: Version: 4.53.3, The environment did not restore to the original 4.57.1.

Required Criteria for Fix

• Track full install provenance/spec, not just normalized name -> version.
• Preserve original install form: -e path, name @ url, name==version
• Restore packages using their original spec format.
• Maintain existing == rollback behavior.
• Replace version-only rollback with spec-aware restore logic.

Summary

• RequirementsManager currently assumes all packages are version-pinned (==).
• This assumption breaks for editable and direct-reference installs, leading to incomplete rollback and persistent environment mutations.

Attempt to Track @ Packages to Validate Rollback Behavior

• Introduced logic to detect and track packages installed via direct references (name @ ...) during the environment snapshot phase. (change and logs - gpt_2.log, pi_0.log) (all their changed/newly installed versions are rolled back here when tracking of "@" is added))

• With this change:
  a) @-based packages are now included in diff computation.
  b) Their original installation spec is preserved instead of being normalized to name==version.

• This means that if the rollback logic is extended to properly handle non-== specs (such as @ and -e installs), these packages can be restored correctly using their original install form — just like ==-pinned packages are handled today