rewriting codebase (#7)

916755e 3 months ago

35.6 kB

Task 18: Legacy Compatibility Audit and Cleanup

Perform a focused cleanup of the GNN4Colliders rewrite now that the ROOT-GNN v1 stack, documentation, performance work, and ONNX export path are in place.

This task should determine exactly which compatibility shims and legacy references are still necessary, remove obsolete migration scaffolding, consolidate remaining compatibility code, and leave the new implementation cleanly separated from the frozen legacy tree.

This task builds on:

Task 8: ROOT-GNN model and transfer/fine-tuning
Task 9: task/loss/metric semantics
Task 11: checkpoint compatibility
Task 12: inference/output compatibility
Task 13: Hydra configuration and CLI
Task 14: DDP/Perlmutter execution
Task 16: documentation and migration closure
Task 17: ONNX export

Before making changes, read:

AGENTS.md
README.md
docs/architecture.md
docs/migration.md
docs/export.md

and any compatibility-specific documentation introduced in earlier tasks.

Then inspect:

src/gnn4colliders/
tests/
configs/
scripts/
legacy/

Pay particular attention to:

src/gnn4colliders/models/root_gnn/
src/gnn4colliders/training/
src/gnn4colliders/inference/
src/gnn4colliders/data/
tests/parity/

Do not delete the legacy source tree in this task unless the repository explicitly no longer needs it for parity/reference purposes.

Goal

Finish the migration boundary cleanly.

The desired architecture is:

new production code
    ↓
clean public APIs
    ↓
small explicit compatibility adapters
    ↓
legacy artifacts / checkpoints / output formats

not:

new production code
    ↓
hidden legacy assumptions everywhere

By the end of this task, it should be clear:

Which legacy behaviors remain supported?
Which compatibility adapters are required?
Which migration-only shims can be removed?
Which legacy concepts no longer exist in public APIs?
Which historical formats are intentionally unsupported?

1. Audit all references to `legacy/`

Search the entire new codebase for:

legacy
root_gnn_dgl
tracking
tracking_info
module.
_orig_mod.
sys.path
dynamic import

Classify every hit.

Use categories such as:

production compatibility
test-only parity
documentation/history
obsolete migration scaffolding
unexpected coupling

Do not remove code simply because it contains the word legacy.

2. No production imports from legacy

The new production package should not import executable logic from:

legacy/

Verify:

src/gnn4colliders/

has no dependency on legacy modules at runtime.

If any exist, replace them with new implementations or explicit compatibility parsing.

Parity tests may still import or execute legacy code.

3. Freeze the compatibility boundary

Identify all remaining compatibility responsibilities.

Likely examples include:

legacy checkpoint loading
legacy state_dict key normalization
legacy classifier/backbone mapping
legacy NPZ output compatibility
legacy metadata/tracking conversion

Keep these responsibilities isolated.

Prefer modules named clearly, such as:

compat/
legacy_checkpoint.py
legacy_outputs.py
legacy_metadata.py

only if multiple compatibility modules genuinely justify a package.

Do not create a broad abstraction layer without need.

4. Consider a dedicated compatibility package

If compatibility code is currently scattered across:

models/
training/
inference/
data/

consider consolidating it under:

src/gnn4colliders/compat/

For example:

src/gnn4colliders/compat/
    __init__.py
    checkpoint.py
    metadata.py
    outputs.py

Use this only if it improves boundaries.

Do not move clean core APIs into compat/.

5. Public API must remain modern

Public APIs should use:

EventMetadata
GraphBatch
sample_id
weight
fold
EdgeNetwork
FineTunedEdgeNetwork

Do not expose public APIs requiring:

tracking[:, 0]
tracking[:, 1]
tracking_info
legacy module/class paths

Legacy structures should appear only at explicit compatibility entry points.

6. Remove positional tracking from new production code

Search for all uses of:

tracking
tracking_info

in production code.

For each:

replace with named metadata if accidentally retained
keep only if it is an explicit compatibility adapter
document why it remains

The normal new pipeline must not construct or propagate a positional tracking tensor.

7. Legacy metadata adapter

If support for old datasets/outputs still requires positional tracking, provide one explicit conversion boundary.

Conceptually:

metadata = EventMetadata.from_legacy_tracking(
    tracking_row,
    ...
)

or equivalent.

This helper should map known positions explicitly:

tracking[0] -> fold
tracking[1] -> weight

Do not propagate the original array beyond the adapter.

8. Reject ambiguous tracking layouts

Do not attempt to infer arbitrary historical tracking layouts.

If a compatibility adapter only understands:

column 0 = fold
column 1 = weight

validate the input shape and fail clearly otherwise.

Do not silently guess.

9. State-dict compatibility audit

Review all compatibility handling for:

module.
_orig_mod.
legacy module names
renamed model components
classifier mappings

Ensure there is one canonical normalization/mapping implementation.

Do not keep duplicate key-remapping code in:

models/
training/checkpoint.py
export/

Reuse one compatibility layer.

10. State-dict normalization should be explicit

Prefer functions such as:

normalize_legacy_state_dict_keys(...)
map_legacy_edge_network_state_dict(...)

with deterministic behavior.

Do not repeatedly use broad:

strict=False

without validating missing/unexpected keys.

11. Remove obsolete Task 8 compatibility helpers

If Task 11 or Task 18 has superseded temporary compatibility helpers introduced during model migration, remove the duplicate helper.

Update imports/tests accordingly.

There should be one supported mapping path.

12. Checkpoint compatibility matrix

Create or update a precise compatibility matrix.

For example:

Artifact	Supported	Notes
new-format checkpoint	yes	full resume
legacy model weights	yes	mapped to new ROOT-GNN
legacy pretrained backbone	yes	fine-tuning
legacy optimizer state	maybe	document exact support
legacy scheduler state	maybe	document exact support
arbitrary historical checkpoint variants	no	unsupported

Use actual implementation status.

13. New checkpoints remain canonical

Do not alter the Task 11 new checkpoint format merely to resemble historical files.

Legacy files adapt into the new representation.

The new checkpoint schema remains the canonical supported format.

14. Legacy checkpoint conversion

If useful, provide an explicit conversion utility:

convert_legacy_checkpoint(...)

that produces a new-format checkpoint.

Only implement this if it simplifies long-term support.

A conversion utility should:

load legacy artifact
normalize model keys
extract supported state
create new checkpoint schema
write explicit compatibility metadata

Do not claim full conversion of unsupported optimizer/trainer state.

15. Conversion should be explicit

Do not silently rewrite a legacy checkpoint on load.

Conversion is a separate user action/API.

Ordinary compatibility loading may continue to work without mutating the source file.

16. Legacy output compatibility audit

Review Task 12 support for historical output fields such as:

scores
labels
tracking_info

Determine whether a real downstream consumer still requires legacy output.

If yes:

keep one explicitly named legacy writer
document it as compatibility-only

If no:

remove the writer
remove stale config options
remove tests that only protect unused historical output

Do not preserve compatibility indefinitely without a use case.

17. Legacy writer naming

If retained, make compatibility obvious.

Prefer:

write_legacy_npz(...)

rather than making legacy behavior the generic:

write_npz(...)

The default writer should use the modern named schema.

18. Modern NPZ remains canonical

The primary output format should retain named fields such as:

sample_id
logits
scores
predictions
labels
weight
fold

as implemented.

Do not revert to positional matrices for legacy convenience.

19. ROOT output compatibility

Audit whether ROOT score writing contains legacy-specific branch naming or behavior.

Separate:

modern output contract

from:

legacy-compatible branch layout

if both are required.

Do not embed hidden legacy names in generic inference logic.

20. Config compatibility audit

Search for any compatibility with old YAML patterns such as:

module:
class:
args:

The new Hydra configuration should not support arbitrary legacy dynamic import syntax.

If a config conversion helper exists, isolate it as an offline migration tool.

Do not put legacy YAML interpretation in the normal CLI path.

21. Legacy config conversion

Only if useful, add a small conversion/reference tool that helps users manually map active legacy configs to new semantic configs.

This tool may extract recognized values.

It should not promise arbitrary automatic translation.

Prefer documentation over complex parsing if only a few active configs exist.

22. Remove stale config aliases

Review Hydra config aliases introduced during migration.

Remove aliases that:

are undocumented
were temporary
duplicate canonical names
are no longer used

Keep aliases only when compatibility value outweighs confusion.

Document retained aliases.

23. CLI compatibility audit

Review the CLI for old naming carried forward merely for migration convenience.

Prefer current commands:

prepare
train
evaluate
predict
export

Do not add old script-name aliases unless actively needed.

Remove stale deprecated flags/config hooks if they have no supported use.

24. Remove `sys.path` migration hacks

Search the entire new package for:

sys.path.insert(...)
sys.path.append(...)

Production package code should not modify import paths.

Remove any migration-era hacks.

Tests should use normal installed-package imports.

25. Import architecture audit

Verify:

gnn4colliders.data
gnn4colliders.features
gnn4colliders.graphs
gnn4colliders.models
gnn4colliders.tasks
gnn4colliders.training
gnn4colliders.inference
gnn4colliders.distributed
gnn4colliders.export

do not depend on arbitrary working-directory layout.

Normal package installation must be sufficient.

26. Remove obsolete global mutable state

Audit new production code for migration-era globals introduced to emulate legacy behavior.

Examples:

global config objects
global dataset state
global device state
global current fold
global checkpoint path

Remove low-risk accidental globals.

Do not perform a major refactor if state is intentional and properly encapsulated.

27. RNG compatibility cleanup

Ensure no compatibility code has reintroduced:

torch.manual_seed(...)
np.random.seed(...)

inside constructors or hidden helpers.

Legacy initialization differences should be handled by:

fixed-weight parity
explicit caller seeding
state_dict compatibility

not hidden global seed mutation.

28. Model architecture cleanup

Review models/root_gnn/ for migration-only structure.

Remove:

unused aliases
temporary wrappers
duplicate legacy class names
dead code paths

only if no public/test/compatibility user remains.

Keep clean modern names canonical.

29. Fine-tuning compatibility cleanup

Ensure transfer learning now uses the explicit Task 8 API.

Remove any remaining code that:

mutates arbitrary model.children()
removes "last layer" by numeric index
depends on legacy module ordering

unless contained entirely in a compatibility state-dict converter.

30. Training compatibility cleanup

Review Trainer/Task integration for legacy terminology such as:

finish function
tracking
test-as-validation

Modern production code should use:

task
metadata
validation
test

Keep legacy terms only in compatibility docs/tests.

31. Loss compatibility audit

Task 9 intentionally preserved active legacy mathematical semantics.

Do not "clean up" the mathematics in Task 18.

Instead ensure compatibility comments/documentation distinguish:

legacy-compatible scientific behavior

from:

legacy implementation structure

The former remains; the latter should be removed where unnecessary.

32. Padding compatibility audit

Review any retained legacy padding modes.

Classify each as:

active required behavior
supported compatibility
dead historical behavior

Remove code for unused modes only if confident no standard/reference config uses them.

Do not remove behavior solely because it looks odd.

33. Cache compatibility audit

Review cache schema handling.

Ensure new code does not silently load old incompatible cache artifacts without validation.

If legacy caches are unsupported, state this explicitly.

If conversion is supported, isolate it.

Prefer rebuilding data caches over maintaining complex cache compatibility unless a real operational need exists.

34. Feature schema compatibility

Ensure:

FEATURE_SCHEMA_VERSION
GRAPH_SCHEMA_VERSION
CACHE_SCHEMA_VERSION

or equivalent are current and documented.

Compatibility adapters should not bypass schema checks casually.

35. Export compatibility audit

Review Task 17 for duplicated compatibility logic.

ONNX export should consume the normalized new model/checkpoint interface.

It should not independently know historical checkpoint key prefixes.

Remove such duplication if present.

36. DDP compatibility audit

New distributed checkpoints should already avoid module. prefixes.

Confirm DDP production code does not rely on legacy prefix handling.

Legacy prefix normalization belongs only to compatibility loaders.

37. Tests: classify parity vs compatibility

Review:

tests/parity/

and other compatibility tests.

Separate mentally and, where helpful, structurally:

scientific parity tests
artifact compatibility tests

Do not remove scientific parity tests merely because migration is complete.

They remain valuable regression tests.

38. Keep core parity tests

Continue protecting:

feature construction
graph topology
edge features
ROOT-GNN forward
fine-tuning forward
loss semantics
metrics

against established legacy behavior.

These are not temporary migration tests if they encode important scientific contracts.

39. Remove obsolete migration tests

Delete tests only when they protect implementation scaffolding that no longer exists.

Examples:

temporary adapter created before final API
duplicate state-dict converter
deprecated config alias with no support commitment

Do not reduce useful behavioral coverage.

40. Legacy executable test dependency

Determine whether parity tests currently require running the entire legacy stack/environment.

If so, consider whether a subset can be converted into frozen golden fixtures.

For example:

legacy input
expected features
expected edge indices
expected edge features
fixed state_dict
expected logits
expected loss

This can reduce dependence on an increasingly old Python/DGL environment.

Do not eliminate direct legacy execution tests if they still provide unique value.

41. Golden fixture strategy

Where appropriate, store small deterministic expected values rather than large artifacts.

Prefer:

tiny tensors
small JSON/NPZ fixtures

Do not commit large ROOT/model binaries solely for compatibility tests.

42. Legacy environment isolation

If some parity tests require the old environment, mark/document them clearly.

They should not make ordinary modern development impossible.

Use the project's existing pytest marker conventions if present.

Do not silently skip critical parity coverage in the canonical validated environment.

43. `legacy/` README

Add or update:

legacy/README.md

explaining:

why legacy exists
which tree is the actual behavioral reference
which abandoned rewrites are not targets
whether legacy code is frozen
whether production code may import it
when it may eventually be removed

Keep this concise.

44. Mark abandoned rewrites clearly

If directories such as:

legacy/physicsnemo/

remain, mark them clearly as:

historical/abandoned rewrite attempt
not parity target
not production implementation

Avoid future developer confusion.

45. Documentation cleanup

Update:

README.md
docs/architecture.md
docs/migration.md
docs/export.md

to reflect the final compatibility boundary.

Remove language suggesting the new rewrite is still transitional if it is now canonical.

46. Migration document status

docs/migration.md should clearly state:

ROOT-GNN rewrite is canonical
legacy remains reference/compatibility source

and list the remaining supported legacy artifacts.

Do not leave "TODO migrate X" entries for completed functionality.

47. Compatibility support policy

Add a concise support policy.

For example:

We preserve:
- active scientific behavior through parity tests
- explicitly documented historical checkpoint formats
- explicitly documented legacy output formats where needed

We do not preserve:
- arbitrary old dynamic-import configs
- undocumented historical cache formats
- every experimental model class

Use actual project decisions.

48. Deprecation language

If any compatibility API is intended for later removal, mark it clearly.

Use:

deprecated compatibility API

only if there is an actual replacement and removal intent.

Do not add noisy runtime warnings everywhere unless helpful.

Documentation may be sufficient.

49. Avoid breaking supported users casually

Do not remove a compatibility behavior that is already documented as supported without:

updating docs
updating tests
explaining the intentional break

This task is cleanup, not arbitrary breaking-change release.

50. Search for stale legacy names

Search for historical class/module names such as:

Edge_Network
Transferred_Learning_Finetuning
GCN.py
tracking
finish

in new production code.

Modern code should generally use:

EdgeNetwork
FineTunedEdgeNetwork
task postprocessing
metadata

Compatibility/documentation references are fine.

51. Search for duplicate behavior

Look for multiple implementations of:

phi wrapping
state_dict normalization
checkpoint reconstruction
score conversion
event weight extraction
sample ID creation

Migration often creates duplicates.

Consolidate only where clearly redundant.

Do not create a giant generic utility module.

52. No `utils.py` dumping ground

Do not solve cleanup by moving unrelated compatibility functions into:

utils.py

Prefer domain-specific modules.

53. Public export cleanup

Review all package __init__.py exports.

Remove public exports for:

migration-only helpers
internal adapters
obsolete aliases

Keep the supported public surface small.

Compatibility helpers may be imported from:

gnn4colliders.compat

if intentionally user-facing.

54. Public API import test

Add or update a test that verifies the intended public API imports.

For example:

from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork
from gnn4colliders.tasks import BinaryClassificationTask
from gnn4colliders.training import Trainer, CheckpointManager
from gnn4colliders.inference import Predictor

Use actual public APIs.

Do not include deprecated internal helpers.

55. Optional compatibility imports

If compatibility is intentionally exposed, make it obvious:

from gnn4colliders.compat import load_legacy_checkpoint

rather than exposing it as a normal checkpoint API with ambiguous semantics.

56. Dependency cleanup

Review pyproject.toml.

Remove dependencies introduced during migration that are no longer used.

Do not remove dependencies based only on static import search if they are required by optional workflows/config.

Verify extras:

root-gnn
onnx/export
dev

or actual current names.

57. Optional-extra boundaries

Ensure:

core install
ROOT-GNN install
ONNX install
development install

remain logically separated.

Legacy compatibility should not force users to install the legacy environment.

58. Legacy dependencies must not leak

The new environment should not require:

Python 3.8
old DGL
old CUDA
legacy ROOT Python bindings

merely to import/use the new production package.

Only dedicated parity/reference workflows may require historical dependencies.

59. Remove unused files

Identify obviously obsolete new-tree files created during earlier tasks.

Examples:

empty compatibility modules
placeholder scripts
superseded config files
unused migration helpers

Delete only after confirming no import/config/test references.

60. Generated artifacts

Remove accidental repository artifacts such as:

__pycache__/
*.egg-info/
profiling traces
temporary checkpoints
temporary NPZ/ONNX outputs

if present.

Update .gitignore where appropriate.

Do not delete legitimate fixtures.

61. Dead-code detection

Use search and test coverage reasoning to identify obvious dead code.

Do not add a new static-analysis dependency solely for this task.

Prefer targeted manual audit.

62. TODO/FIXME audit

Search:

TODO
FIXME
HACK
XXX
legacy
temporary
migration

in production code.

Resolve or reclassify each relevant migration-era note.

Do not leave stale comments claiming something is temporary after it became the final implementation.

63. Comments should explain compatibility

Keep compatibility comments focused on why behavior exists.

For example:

Legacy checkpoints may contain the `module.` prefix from DDP.
Normalize it at this boundary so the canonical model state remains prefix-free.

Avoid comments that narrate obvious code.

64. Error messages

Compatibility failures should explain what is unsupported.

For example:

Unsupported legacy checkpoint layout: expected model_state_dict and epoch fields.

Prefer actionable errors over raw KeyError/shape errors.

Do not catch every exception indiscriminately.

65. Compatibility warnings

When loading supported-but-legacy artifacts, consider a concise informational warning/log.

For example:

Loading legacy ROOT-GNN checkpoint through compatibility adapter.

Do not spam warnings per tensor/key.

Make warning behavior testable if added.

66. Unsupported legacy artifacts

Explicitly reject known unsupported formats rather than partially loading them.

Examples may include:

experimental model families
unknown tracking layout
obsolete cache schema
unrecognized checkpoint architecture

Fail early.

67. Experimental legacy models

Confirm the new production code does not accidentally expose unsupported historical models such as:

GCN_global
GCN_global_2way
attention variants
MultiModel
Clustering

unless they were intentionally migrated in a later task.

Document them as unsupported legacy experiments.

68. Legacy transfer variants

Likewise, retain only the active transfer/fine-tuning behavior migrated in Task 8.

Do not keep compatibility wrappers for unused historical transfer classes without evidence they are needed.

69. Legacy config examples

Remove obsolete old YAML examples from the main documentation.

If valuable historically, keep them under:

legacy/

or migration docs.

Normal users should see semantic Hydra configs.

70. Re-run active reference configs

After cleanup, ensure the canonical reference workflows still compose/run:

multiclass pretraining
binary training
pretrained fine-tuning
resume
evaluate
predict
export

Cleanup must not break the v1 stack.

71. Checkpoint round-trip regression

Run Task 11 checkpoint tests.

Ensure compatibility cleanup does not affect:

new-format save/load
resume
best/latest selection
weight-only loading

72. Legacy checkpoint regression

Run explicit legacy checkpoint compatibility tests.

Verify all formats still documented as supported remain supported.

Do not weaken expectations merely to simplify adapters.

73. Fine-tuning regression

Verify:

legacy pretrained checkpoint
    ->
compatibility loader
    ->
new backbone
    ->
FineTunedEdgeNetwork

still works for the supported legacy artifact.

This is one of the most valuable compatibility workflows.

74. Inference compatibility regression

If modern inference supports loading historical model weights, test:

legacy checkpoint
    ->
new Predictor
    ->
modern PredictionResult

The output should use the modern schema unless the user explicitly requests a legacy writer.

75. ONNX compatibility regression

If a supported legacy checkpoint can be normalized into a new model, verify it can still reach the Task 17 export path.

Do not make ONNX exporter directly consume legacy layouts.

The flow should be:

legacy checkpoint
    ->
compatibility normalization
    ->
new model
    ->
ONNX export

76. No legacy logic inside ONNX adapter

Verify Task 17's export adapter has no:

legacy key normalization
tracking conversion
old class-name handling

Those belong upstream.

77. DDP regression

Run Task 14 tests.

Cleanup of state-dict prefix handling must not break:

new DDP checkpoint save
single-process load
distributed resume

New DDP behavior and legacy DDP compatibility are separate concepts.

78. Performance regression

Do not intentionally degrade Task 15 optimized hot paths just to centralize compatibility code.

Compatibility adapters should generally operate at:

load
conversion
output

boundaries, not inside every training batch.

79. Compatibility code should not be in hot loops

Ensure no legacy conversion runs:

per node
per edge
per batch

during normal modern training.

Convert once at ingestion/load boundaries.

80. Architecture-neutral audit

Use cleanup to verify shared infrastructure remains future-model friendly.

Search for ROOT-GNN compatibility assumptions inside:

data/
features/
tasks/
training/
inference/
distributed/

Move graph/model-specific compatibility toward:

compat/
graphs/
models/root_gnn/

where appropriate.

Avoid major refactors.

81. Future ROOT-Transformer must not inherit legacy tracking

Ensure no shared API requires:

tracking
GraphSample
DGLGraph

for all models.

Shared EventSample/metadata should remain architecture-neutral.

Document any remaining coupling as technical debt.

82. Test fixture cleanup

Review test fixtures for duplicate or stale legacy data representations.

Prefer reusable tiny fixtures for:

EventSample
GraphSample
GraphBatch
state_dict

Keep explicit legacy fixtures only where testing compatibility.

83. Avoid huge fixture binaries

Do not commit real production checkpoints/datasets.

Use tiny deterministic synthetic fixtures.

If a legacy checkpoint fixture is needed, construct a minimal compatible dictionary in tests.

84. Compatibility tests naming

Use clear test names such as:

test_legacy_checkpoint_module_prefix
test_legacy_checkpoint_orig_mod_prefix
test_legacy_tracking_to_metadata
test_legacy_npz_writer

Do not hide compatibility behavior inside generic tests.

85. Deprecation tests

If an API is formally deprecated, add a minimal test for the warning/behavior.

Do not test exact warning prose too rigidly.

86. Documentation: compatibility guide

If compatibility behavior is substantial, add:

docs/compatibility.md

Otherwise keep it in docs/migration.md.

Document:

supported historical checkpoints
supported output compatibility
unsupported artifacts
conversion examples

Keep it practical.

87. Example legacy checkpoint conversion

If a conversion API is implemented, document a small example.

Conceptually:

uv run python -m ... convert-legacy-checkpoint ...

or a Python API.

Do not add a CLI command unless conversion is a real user workflow.

A documented Python utility may be enough.

88. No broad legacy CLI

Do not add:

gnn4colliders legacy ...

as a generic compatibility subsystem.

Keep compatibility focused.

89. Migration closure statement

Update migration docs to state that production development should now target:

src/gnn4colliders/

and not modify the frozen legacy implementation except when maintaining parity fixtures/documentation.

90. Legacy deletion criteria

Document what must be true before legacy/ can eventually be removed.

For example:

all required scientific contracts have stable golden fixtures
historical checkpoint compatibility no longer requires direct source inspection
no active parity test executes legacy code
repository policy approves archival/removal

Do not delete it in this task unless explicitly instructed.

91. Validate documentation links

Ensure README/docs references to:

compatibility
migration
legacy
checkpoint conversion

point to real files/sections.

Avoid broken links after moving docs.

92. Validation

Run searches such as:

grep -R "tracking\\[" src/gnn4colliders
grep -R "tracking_info" src/gnn4colliders
grep -R "sys.path" src/gnn4colliders
grep -R "legacy/" src/gnn4colliders

Use a more appropriate search tool if available.

Review every hit.

Run focused compatibility tests:

uv run pytest tests/parity -v

Run checkpoint tests:

uv run pytest tests/unit/training -v

Run inference tests:

uv run pytest tests/unit/inference -v

Run export tests:

uv run pytest tests/unit/export -v

Run integration tests:

uv run pytest tests/integration -v

Run distributed tests:

uv run pytest tests/unit/distributed -v
uv run pytest tests/integration -k distributed -v

Run the full suite:

uv run pytest

Run lint/format:

uv run ruff check .
uv run ruff format --check .

Validate public CLI workflows:

uv run gnn4colliders --help
uv run gnn4colliders prepare --help
uv run gnn4colliders train --help
uv run gnn4colliders evaluate --help
uv run gnn4colliders predict --help
uv run gnn4colliders export --help

Inspect:

git status
git diff

Verify:

production code does not import executable legacy modules
positional tracking is absent from modern APIs
compatibility logic is isolated
new checkpoints remain canonical
modern output schemas remain canonical
no duplicate state-dict mapping remains
no old dynamic-import YAML path exists in normal execution
scientific parity tests remain intact
no supported legacy artifact was broken unintentionally
no major new functionality was added
no unrelated refactor is included

Completion criteria

Task 18 is complete when:

Every production reference to legacy concepts has been audited.
Production code does not import executable logic from legacy/.
Positional tracking arrays are absent from modern public APIs.
Legacy tracking conversion, if retained, exists only at an explicit compatibility boundary.
Legacy checkpoint key normalization has one canonical implementation.
Duplicate state-dict mapping helpers have been removed.
New checkpoint format remains canonical.
Supported legacy checkpoint workflows still pass.
Unsupported historical checkpoint behavior is documented clearly.
Resume and pretrained transfer remain distinct.
Modern NPZ output remains canonical.
Legacy NPZ/output behavior is explicitly compatibility-only or removed.
Old dynamic-import YAML is absent from normal CLI/config paths.
sys.path migration hacks are absent from production code.
Obsolete migration globals/helpers are removed.
ROOT-GNN production model code uses modern APIs/names.
Task/training code contains no accidental legacy positional semantics.
Cache/schema compatibility behavior is explicit.
Task 17 export uses normalized new models rather than legacy logic.
DDP new-format state handling remains unaffected.
Scientific parity tests remain intact.
Obsolete migration-only tests are removed only where justified.
legacy/README.md clearly describes the frozen legacy tree.
Abandoned historical rewrites are clearly marked.
Documentation accurately describes supported compatibility.
Public exports are cleaned up.
Unused dependencies/files introduced during migration are removed where safe.
Full modern pretraining/fine-tuning/inference/export workflows still work.
Full tests pass.
Lint/format checks pass.
Legacy source remains frozen.

Completion report

Report:

files created
files modified
files removed
legacy-reference audit findings
positional tracking audit findings
compatibility package/module structure
legacy metadata conversion behavior
checkpoint compatibility consolidation
state-dict normalization behavior
legacy checkpoint formats supported
legacy checkpoint formats unsupported
legacy optimizer/scheduler resume status
legacy output compatibility status
config compatibility cleanup
CLI compatibility cleanup
sys.path/import cleanup
model migration-shim cleanup
training/task cleanup
cache/schema compatibility findings
export compatibility cleanup
DDP compatibility regression results
parity-test cleanup
golden-fixture changes
legacy executable-test dependency status
public API cleanup
dependency cleanup
generated/dead-file cleanup
TODO/FIXME audit results
documentation updates
remaining compatibility debt
criteria for eventual legacy/ removal
validation commands and results

After validation succeeds, create one Git commit containing only Task 18 changes.

Use:

refactor: consolidate legacy compatibility boundaries

Before committing, inspect the final diff and ensure no unrelated scientific or architectural changes are included.

Task 18: Legacy Compatibility Audit and Cleanup

Goal

1. Audit all references to legacy/

2. No production imports from legacy

3. Freeze the compatibility boundary

4. Consider a dedicated compatibility package

5. Public API must remain modern

6. Remove positional tracking from new production code

7. Legacy metadata adapter

8. Reject ambiguous tracking layouts

9. State-dict compatibility audit

10. State-dict normalization should be explicit

11. Remove obsolete Task 8 compatibility helpers

12. Checkpoint compatibility matrix

13. New checkpoints remain canonical

14. Legacy checkpoint conversion

15. Conversion should be explicit

16. Legacy output compatibility audit

17. Legacy writer naming

18. Modern NPZ remains canonical

19. ROOT output compatibility

20. Config compatibility audit

21. Legacy config conversion

22. Remove stale config aliases

23. CLI compatibility audit

24. Remove sys.path migration hacks

25. Import architecture audit

26. Remove obsolete global mutable state

27. RNG compatibility cleanup

28. Model architecture cleanup

29. Fine-tuning compatibility cleanup

30. Training compatibility cleanup

31. Loss compatibility audit

32. Padding compatibility audit

33. Cache compatibility audit

34. Feature schema compatibility

35. Export compatibility audit

36. DDP compatibility audit

37. Tests: classify parity vs compatibility

38. Keep core parity tests

39. Remove obsolete migration tests

40. Legacy executable test dependency

41. Golden fixture strategy

42. Legacy environment isolation

43. legacy/ README

44. Mark abandoned rewrites clearly

45. Documentation cleanup

46. Migration document status

47. Compatibility support policy

48. Deprecation language

49. Avoid breaking supported users casually

50. Search for stale legacy names

51. Search for duplicate behavior

52. No utils.py dumping ground

53. Public export cleanup

54. Public API import test

55. Optional compatibility imports

56. Dependency cleanup

57. Optional-extra boundaries

58. Legacy dependencies must not leak

59. Remove unused files

60. Generated artifacts

61. Dead-code detection

62. TODO/FIXME audit

63. Comments should explain compatibility

64. Error messages

65. Compatibility warnings

66. Unsupported legacy artifacts

67. Experimental legacy models

68. Legacy transfer variants

69. Legacy config examples

70. Re-run active reference configs

71. Checkpoint round-trip regression

72. Legacy checkpoint regression

73. Fine-tuning regression

74. Inference compatibility regression

75. ONNX compatibility regression

76. No legacy logic inside ONNX adapter

77. DDP regression

78. Performance regression

1. Audit all references to `legacy/`

24. Remove `sys.path` migration hacks

43. `legacy/` README

52. No `utils.py` dumping ground