test(anaconda-models): validate local cache reuse in Metaflow [AIC-4186]

[inhaz1]

Summary

Adds a Metaflow E2E flow to validate task-local cache reuse for @anaconda_models.

What this tests

• First pull downloads the model into a clean local cache directory.
• Validates the expected model name, format, and quantization.
• Confirms downloaded file status and size.
• Creates a new model handle in the same Metaflow task.
• Verifies the second pull reuses the cached model and reports skipped.
• Confirms cached file count and per-file statuses.
• Verifies the cached file path, size, and modification time remain unchanged.

Scope

This covers cache reuse within a single Metaflow task/process. Shared cache behavior across separate tasks or pods is out of scope.

Test

python models/local_cache_model_flow.py --environment=fast-bakery run --with kubernetes

Jira: AIC-4186

[Copilot]

🟢 Approval recommended

The focused test fully covers the cache-reuse behavior described in the PR.

Pull request overview

Adds a Metaflow E2E flow validating task-local model cache reuse.

Changes:

• Verifies cold download metadata and file state.
• Verifies a second handle reuses the unchanged cached file.

File summaries

File	Description
flows/models/local_cache_model_flow.py	Adds the local-cache reuse E2E flow.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 0
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟢 Approval recommended

The focused E2E flow correctly validates the stated cache-reuse behavior.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[vvemulapalli11]

@inhaz1 do we need to add access_denied_reason guard in this scenario too?

[vvemulapalli11]

@inhaz1 since this is a shared file system, removing will impact other concurrent tests

1. CACHE_ROOT is a fixed path, and the test deletes it at the start. If two runs of this flow happen at the same time, one run deletes the model the other just downloaded. The second pull then says "downloaded" instead of "skipped" and the test fails, or the delete itself crashes with FileNotFoundError. (Only this flow sets temp_dir_root, so the other tests are not affected.)

We cannot put the run ID in the path, because the decorator reads it when the file is imported, before a run ID exists. But the flow file is imported again inside the task, so a unique value at the top of the file is picked up by both the decorator and the step.

2. The test decides the cache was empty by trusting download_status == "downloaded" — but that is the thing we are testing. It is also not proof: if a leftover file is the wrong size, the client downloads again anyway. So we should check the file is gone ourselves.

Both fixes below. Instead of deleting the whole folder, delete only this model's file. .model() does not download anything (pull=False by default), and for a single file the client only checks that the file exists and has the right size — so deleting that one file is enough.

# One folder per task process, so runs cannot delete each other's model.
CACHE_ROOT = os.path.join(
    tempfile.gettempdir(),
    f"obp-anaconda-model-local-cache-test-{os.getpid()}-{uuid.uuid4().hex[:8]}",
)


class LocalCacheModelFlow(FlowSpec):
    @anaconda_models(temp_dir_root=CACHE_ROOT)
    @step
    def start(self):
        try:
            self._assert_cache_reuse()
        finally:
            # Don't leave a large model file behind in /tmp.
            shutil.rmtree(CACHE_ROOT, ignore_errors=True)

        self.next(self.end)

    def _assert_cache_reuse(self):
        # Nothing is downloaded here, so the file path is known but not created.
        cold_model = self.anaconda_models.model(...)

        # ... existing identity assertions unchanged ...

        assert cold_model.path.startswith(CACHE_ROOT + os.sep), (
            f"Expected the model under {CACHE_ROOT}, got {cold_model.path}"
        )
        # For a collection, .path is a folder, not a file.
        assert not cold_model.is_collection, (
            f"This setup expects a single model file: {cold_model.path}"
        )

        # Remove only this model, then check for ourselves that it is gone.
        if os.path.exists(cold_model.path):
            os.remove(cold_model.path)
        assert not os.path.exists(cold_model.path), (
            f"Cache is not empty, model file still there: {cold_model.path}"
        )

        cold_model.pull()

        # ... rest of the existing cold/warm assertions 

[vvemulapalli11]

@inhaz1 I am not sure here, please double check
The current assertion can fail when the catalog doesn’t provide size_bytes, since .size will be None. This causes the test to stop before reaching the warm-pull assertions, even though caching is working correctly. Could we use the same minimum_size_mb check as the sibling flow instead?

[vvemulapalli11]

feedback comments added

[Copilot]

🟡 Changes recommended

The promised inode-preservation validation is missing.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Balanced

[inhaz1]

not needed
The existing assertions sufficiently validate cache reuse, so I removed this claim from the PR description.

[vvemulapalli11]

Review: local_cache_model_flow.py

I traced the assertions in this flow against the @anaconda_models decorator and AnacondaModel/AnacondaModelClient implementation in flows/.venv/.../anaconda_models/{decorator,client}.py. The core logic is sound:

• _pull_single() sets download_status = "downloaded" on a cache miss and "skipped" on a hit (via _is_present()), matching the assertions at lines 57 and 91.
• On a "skipped" pull, the existing file is never rewritten, so inode/mtime/size are preserved — matching the assertions at lines 105–113.
• root_dir (here CACHE_ROOT) is passed straight through from temp_dir_root, and the destination path is deterministic from model name + quant_method + format, so a fresh model() handle with the same filters resolves to the same path — matching the warm_model.path == cold_path assertion.

A few things worth considering before/after merge:

1. assert warm_model is not cold_model (line 84) is vacuous. AnacondaModelClient.model() always constructs a brand-new AnacondaModel instance, so this identity check can never fail — it doesn't actually verify anything about handle independence. Either drop it or replace it with something that demonstrates state isn't shared (e.g., asserting warm_model.files is not cold_model.files, or that mutating one doesn't affect the other).

2. Fixed CACHE_ROOT (/tmp/obp-anaconda-model-local-cache-test) could collide across concurrent runs on the same node. The comment at line 32 only covers task retries reusing the same path. Since this is a shared literal path under /tmp, two run --with kubernetes executions landing on the same node/container (e.g., concurrent CI runs, or local testing outside k8s) could race on shutil.rmtree/downloads. Not necessarily a blocker if each run always gets an isolated pod, but worth a comment confirming that assumption or namespacing the path with os.getpid()/a run id to be safe.

3. No cleanup of the downloaded file/CACHE_ROOT. Other flows in this directory don't clean up either, so this is consistent, but since this flow explicitly creates/manages a fixed cache directory (rather than a tempfile.mkdtemp()-owned one), it might be nice to remove it in the end step for hygiene, especially if the flow is ever run locally without --with kubernetes.

None of these are correctness issues with the cache-reuse behavior being tested — the test logic itself is valid and mirrors the sibling flows' style (download_gguf_model_flow.py). Items 1–2 are minor robustness/clarity suggestions.

[vvemulapalli11]

LG!