Activity

1. pszemraj

   added

   enhancementNew feature or requestNew feature or request

   on Aug 24, 2026
2. pszemraj commented on Sep 23, 2026

   

   pszemraj

   on Sep 23, 2026

   OwnerAuthor

   More actions

   Splitting codedupes/semantic.py into a Package

   codedupes/semantic.py embeds code units, manages a process-wide model cache, fingerprints local model directories, resolves cache and vector-space identity, runs an OOM-recovery encode ladder, and implements duplicate scanning and query search. The recommendation is to convert it into a codedupes/semantic/ package of thirteen modules arranged as a strict import DAG, with every piece of process-wide mutable model state in exactly one module (model_registry.py). The split is a pure move with a single signature change, and each migration step ships independently behind a frozen-API test.

   Two structural choices carry most of the design. First, all dtype, fast-math, and write-coherence policy lives below the model registry, because the cache-variant builder and the registry both need it. Placing that policy inside the registry creates an import cycle. Second, cache addressing stays in the same module as vector-space identity, because the two derive the same variant string and one suffix is already hand-duplicated between them.

   Contents

   • Target Layout
   • Import Order
   • Placement Decisions
   • Invariants to Preserve
   • Compatibility Surface
   • Migration Sequence
   • Rejected Alternatives
   • Follow-up Work
   • Coarser Variant

   Target Layout

   codedupes/semantic/
     __init__.py        # facade: explicit re-exports + __all__; no logic
     runtime.py         # _configure_semantic_runtime_env, _safe_package_version,
                        # get_semantic_runtime_versions, _require_dependency, _check_semantic_dependencies
     errors.py          # SemanticBackendError, _is_known_semantic_backend_error,
                        # _wrap_semantic_backend_error, _classify_oom_device
     vectors.py         # numpy only: InvalidEmbeddingError, canonicalize_embeddings,
                        # _validate_precomputed_embeddings, _assemble_cached_matrix,
                        # _select_cache_miss_indices, _PRECOMPUTED_VALIDATION_BLOCK_ROWS
     encoding.py        # EncodeRoute, EncodePlan, EMBEDDINGGEMMA_* prompt tables, _resolve_encode_plan,
                        # resolve_encode_plan, _select_encode_fn, _prepare_embedding_text,
                        # _prepare_search_document, _append_context_truncation_diagnostics
     local_model.py     # manifest memo + lock, manifest load/store, _hash_file_content,
                        # _fingerprint_local_model_dir(_cached), walk-scope global + context manager,
                        # _remember_local_model_fingerprint_in_scope, _validate_local_model_directory
     revisions.py       # _resolve_profile_default, _resolve_model_revision, _resolve_trust_remote_code,
                        # _resolve_hf_cached_revision, _is_hf_commit_hash, _resolve_load_revision(_for_cache_policy),
                        # _resolve_revision_for_cache, _get_loaded_model_commit_hash, _revision_is_mutable_label
     device_policy.py   # _resolve_semantic_device_request, validate_explicit_device_request,
                        # _prepare_semantic_device, _warned_mlx_mps_contention, _coerce_device_name,
                        # _resolve_model_dtype, _model_parameter_dtype, _dtype_variant_for,
                        # _mps_fast_math_variant, _fast_math_write_allowed, _dtype_coherence_broken,
                        # _cache_write_allowed
     identity.py        # SearchDocumentMode, _validate_search_document_mode, EmbeddingSpaceIdentity,
                        # EMBEDDING_PIPELINE_SCHEMA, _embedding_runtime_fingerprint, _cache_variant_for,
                        # _build_embedding_space_identity, _CodeEncodeContext, _resolve_code_encode_context,
                        # resolve_embedding_space_identity, _require_current_embedding_space,
                        # embedding_cache_keys_for_units, _embedding_cache_namespace, _prepare_cache_context,
                        # _MAX_CACHED_QUERY_KEYS
     model_registry.py  # every _model* global, _model_lock, _warned_cpu_fallback_reuse,
                        # _get_model_unlocked, get_model, _clear_model_cache_unlocked, clear_model_cache,
                        # _set_model_execution_device, _get_effective_model_device, _move_model_to_cpu,
                        # _confirm_cache_revision_after_load
     inference.py       # _encode_texts, _encode_with_retries
     embed.py           # ProgressMode, PROGRESS_BAR_MIN_INPUTS, _should_show_progress, EmbeddingRunStats,
                        # _reset/_record_embedding_run_stats, _compute_embeddings_unlocked,
                        # compute_embeddings_with_identity, compute_embeddings
     search.py          # QueryExecution, resolve_search_threshold, _find_similar_to_query_unlocked,
                        # find_similar_to_query
     duplicates.py      # _PAIRWISE_SCAN_BLOCK_SIZE, find_semantic_duplicates,
                        # run_semantic_analysis_with_identity, run_semantic_analysis
   

   get_code_unit_statement_count leaves the package entirely and moves next to CodeUnit in codedupes/models.py, since it has nothing to do with embeddings.

   Import Order

   An arrow A --> B means B imports A. Dependencies sit above their dependents, and no module imports anything below it. Sibling modules import each other with relative imports (from .identity import ...). No submodule imports from the package __init__, which becomes a re-export layer and would otherwise create cycles during package initialization.

   flowchart TD
     runtime --> errors
     errors --> local_model
     local_model --> revisions
     runtime --> device_policy
     errors --> device_policy
     revisions --> identity
     device_policy --> identity
     encoding --> identity
     revisions --> model_registry
     device_policy --> model_registry
     model_registry --> inference
     vectors --> inference
     identity --> embed
     inference --> embed
     identity --> search
     inference --> search
     embed --> duplicates
   

   Loading

   Transitive and leaf edges are omitted for readability; vectors and encoding are leaves that anything may import. identity and model_registry are siblings: neither imports the other, and the orchestration modules (embed, search) are the first to import both. The constraint that keeps the graph acyclic is that nothing in revisions, device_policy, encoding, or identity may import model_registry.

   Placement Decisions

   Dtype, fast-math, and write-coherence policy sit in device_policy.py, below the registry. _cache_variant_for (identity) needs _dtype_variant_for, which needs _resolve_model_dtype. The registry's load path also needs _resolve_model_dtype, and the coherence checks need _dtype_variant_for and _mps_fast_math_variant. If these functions live in the registry, then identity imports registry and registry imports identity. With them in device_policy.py, both depend downward on one policy module, and the registry holds only state and the load/release lifecycle.

   _cache_write_allowed takes the execution device as a parameter. It is the only policy function that reads registry state, via _get_effective_model_device. Passing the device in removes the upward dependency. This is the only signature change in the split, and it has two call sites: the corpus write in embed.py and the query write in search.py.

    def _cache_write_allowed(
        profile: SemanticModelProfile,
        device: str | None,
        mps_fallback: bool | None,
        resolved_device: str,
        model: object,
   +    execution_device: str,
    ) -> bool:
   -    execution_device = _get_effective_model_device(model, resolved_device)
        if not _fast_math_write_allowed(device, execution_device):
            return False

   Callers pass _get_effective_model_device(model, resolved_device). They must compute it at the call site, after encoding, so that a mid-encode CPU fallback is still reflected. That timing is the reason the function re-reads the device today.

   _confirm_cache_revision_after_load lives in model_registry.py. It is revision logic, but its local-directory branch takes _model_lock and compares model is _model to return _model_local_fingerprint. Placing it in revisions.py or identity.py would force those modules to import the registry.

   _classify_oom_device lives in errors.py. Both the model load path (registry) and the encode ladder (inference) classify OOMs. Placing it in inference.py would make the registry import inference, which imports the registry.

   runtime.py sits below errors.py. _wrap_semantic_backend_error embeds get_semantic_runtime_versions() in its message, so errors depend on runtime, not the reverse.

   Cache addressing stays in identity.py. EmbeddingSpaceIdentity.runtime_variant is literally _cache_variant_for(...), and cache keys are derived from the same variant. The contextual-document suffix is currently written out twice: once in _compute_embeddings_unlocked through variant_suffix into _prepare_cache_context, and once inline in embedding_cache_keys_for_units. If the two copies drift, manifest keys silently stop matching the vectors the cache holds. Keep both call sites in one module, and collapse them into one helper during the move:

   def _with_search_document(variant: str, search_document: SearchDocumentMode) -> str:
       # Single definition of the contextual-document key suffix; used by
       # _prepare_cache_context and embedding_cache_keys_for_units.
       if search_document == "contextual":
           return f"{variant}\x00search_document=contextual"
       return variant

   The query-only query_provenance=2 suffix stays inside _prepare_cache_context, because it has a single writer.

   The scan module is named duplicates.py, not pairs.py. The code already imports ordered_pair_key from the top-level codedupes/pairs.py. A second pairs module one level down makes greps and tracebacks ambiguous.

   Invariants to Preserve

   State each of these in the docstring of the module that owns it, so the constraint travels with the code.

   One lock object, reentrant for the right reason. _model_lock is defined only in model_registry.py, and embed.py and search.py import that object. It is an RLock because the top-level entry points (compute_embeddings_with_identity, find_similar_to_query) hold it for the whole operation, and inside that operation get_model() and _confirm_cache_revision_after_load acquire it again. The faithful-CPU restart calls _compute_embeddings_unlocked directly and does not re-acquire the lock, so it is not the reason. Documenting the wrong reason invites someone to "fix" the restart and then downgrade the lock.

   Warning flags live beside their log_warning_once(globals(), ...) calls. log_warning_once reads and writes the flag by name in the mapping it is handed, which is the calling module's globals. _warned_mlx_mps_contention moves with _prepare_semantic_device into device_policy.py. _warned_cpu_fallback_reuse stays in model_registry.py, where _get_model_unlocked sets it and _clear_model_cache_unlocked resets it.

   The fingerprint walk-scope global stays unlocked and private to local_model.py. It is safe without its own lock only because it is touched exclusively while the caller holds _model_lock. Only the context manager and the two accessor functions may reference it.

   Model release order in _clear_model_cache_unlocked is load-bearing. The function drops every reference, deletes the local, and only then clears the device cache. It deliberately does not migrate the model to CPU first. Do not "tidy" this during the move.

   The model-free warm path stays import-free. A fully cached corpus returns without importing torch, and _configure_semantic_runtime_env must run before torch is first imported. New modules may import torch, sentence_transformers, and transformers only inside functions, as _resolve_model_dtype, _move_model_to_cpu, and _get_model_unlocked do today. A top-level import added for type hints in any new module silently breaks both contracts.

   Compatibility Surface

   The package __init__.py is the public API and re-exports explicitly. Build the export list from actual usage, found by grepping the repository and tests for from codedupes.semantic import and semantic\.\w+. The list must include private names used outside the module. _prepare_search_document, for example, has no caller inside semantic.py, so something external imports it. Declare the list in __all__, and do not use star imports, because star imports skip underscore names.

   Monkeypatches must target the module where the name is looked up at call time. After a function moves, a test that patches the facade binding no longer affects callers inside the defining module. Such tests keep passing while testing nothing. This was verified directly: with caller() and helper() in sem/impl.py and re-exported from sem/__init__.py, setting sem.helper leaves caller() returning the real value, while setting sem.impl.helper takes effect. Every patch of get_model, _resolve_hf_cached_revision, _get_loaded_model_commit_hash, _fingerprint_local_model_dir, the _model* globals, and similar names must be retargeted in the same commit that moves the name's callers.

   Re-exported globals are snapshots. from .model_registry import _model binds None into the facade permanently. Any code outside the registry that reads semantic._model or semantic._model_execution_device must read model_registry._model directly.

   No semantic.py shim alongside the package. When codedupes/semantic.py and codedupes/semantic/__init__.py both exist, Python imports the package (verified on 3.12). A leftover shim file is dead code that looks authoritative. Use git mv to turn the file into __init__.py so history follows.

   Logger names change harmlessly. Loggers become codedupes.semantic.embed, codedupes.semantic.model_registry, and so on. They are children of codedupes.semantic, so configuration on the old name still applies through propagation.

   Migration Sequence

   Each step is one commit that passes the full test suite. Code not yet moved stays in __init__.py and imports moved names from their new module. Order is by risk, not strictly by the DAG: stateless leaves go first, and the singleton move goes alone.

   1. Rename and add guards. Run git mv codedupes/semantic.py codedupes/semantic/__init__.py. Add two guard tests: one asserts the frozen export set is importable from codedupes.semantic, and one asserts that importing codedupes.semantic loads no heavy modules beyond the pre-split baseline.
   2. Stateless leaves. Move runtime.py, errors.py, vectors.py, encoding.py, and the scan half of duplicates.py (find_semantic_duplicates and its block size). None of these touch shared state.
   3. Fingerprinting and revisions. Move local_model.py and revisions.py together, because _resolve_revision_for_cache depends on the call-scoped fingerprint memo.
   4. Device policy. Move device_policy.py, apply the _cache_write_allowed signature change, and move the MLX warning flag.
   5. Identity. Move identity.py, including cache addressing, and introduce _with_search_document.
   6. Registry. Move model_registry.py alone. Retarget every test patch and external read of registry names in this commit. This is the highest-risk step.
   7. Orchestration. Move inference.py, embed.py, and search.py, then move run_semantic_analysis* into duplicates.py.
   8. Facade. Reduce __init__.py to the explicit re-export list and __all__.

   The guard tests from step 1 are small; this sketch is illustrative:

   import subprocess
   import sys
   
   import codedupes.semantic as semantic
   
   FROZEN_EXPORTS = {"compute_embeddings", "find_semantic_duplicates"}  # extend from the usage grep
   
   def test_export_surface_is_stable():
       missing = {name for name in FROZEN_EXPORTS if not hasattr(semantic, name)}
       assert not missing
   
   def test_import_stays_model_free():
       code = "import sys, codedupes.semantic; print(sorted(m for m in sys.modules if m.split('.')[0] in {'torch', 'transformers', 'sentence_transformers'}))"
       out = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True).stdout
       assert out.strip() == "[]"  # or the recorded pre-split baseline if codedupes.devices imports torch eagerly

   Rejected Alternatives

   Grouping dtype and coherence helpers with the model singleton. The grouping feels natural because those helpers inspect the loaded model. It produces the registry <-> identity cycle described under Placement Decisions and bloats the one module whose move is genuinely risky.

   Moving _prepare_semantic_device into a runtime module while its warning flag goes to the registry. The flag must sit in the same module as its log_warning_once(globals(), ...) call. Splitting them breaks the warn-once behavior, because the call reads and writes a flag that no longer exists in its module.

   A separate cache_addressing.py. This separates the two call sites that share the contextual suffix, which is exactly the lockstep hazard this split should remove.

   A semantic.py re-export shim next to the new package. Python resolves the package first, so the shim would never load.

   Follow-up Work

   The split leaves _compute_embeddings_unlocked at roughly 350 lines and _find_similar_to_query_unlocked at roughly 300. Two extractions address most of that size.

   Cache-session extraction. Both functions repeat the "load model -> confirm revision -> re-key -> re-lookup" sequence and the mutable-branch provenance reconciliation (confirm_source_commit, snapshot source-commit comparison, bypass when the loaded commit is unknown). Moving these into a small cache-session object in a provenance.py would roughly halve both functions and give the "one matrix never mixes two checkpoints" invariant a single implementation.

   Allocator guard helper. Four paths skip _prepare_semantic_device and must call restore_mps_memory_fraction_if_managed() by hand: the empty corpus and full warm hit in _compute_embeddings_unlocked, and the empty corpus and warm query hit in _find_similar_to_query_unlocked. A no_allocator_work() helper would encode that contract once.

   Warning-flag storage (optional). Once each flag is colocated with its call, the flags could move to a per-module _WARNED: dict[str, bool] passed to log_warning_once instead of globals(). This works only if log_warning_once touches nothing but item get/set on its first argument, so check its implementation in embedding_cache.py first.

   Coarser Variant

   If thirteen modules is more granularity than the team wants, merge along the same DAG without changing any placement rule. This yields six modules:

   Module	Contains	Imports from
   foundation.py	runtime, errors, vectors, encoding	third-party only
   policy.py	local_model, revisions, device_policy, identity	foundation
   model_registry.py	unchanged	foundation, policy
   embed.py	inference + embed	everything above
   search.py	unchanged	everything above
   duplicates.py	unchanged	foundation, embed

   The registry stays alone in this variant as well. The point of the split is that the only module holding process-wide model state is small enough to review in one sitting.