Update to Model Hub integration by r-sarma · Pull Request #502 · interTwin-eu/itwinai

r-sarma · 2026-07-21T13:12:02Z

This PR is an update to the RI-SCALE Model Hub integration with itwinai, which was a CLI-based implementation.

Now the AI Model Hub is added as a feature, allowing the abstraction of the backend, when pulling models. Users are able to define this already in the configuration file.
The model pulling API is also added, which allows the users to pull models from the Model Hub and launch inference.

okrochak

I've left some comments about the modelhub-related code organization. I've tried to run the tutorial, also with my own .env file, but didn't manage to (see the last comment).

okrochak · 2026-08-05T07:37:26Z

GitHub preview shows me a lot of text being underlined, which I think is not so good for readability.

okrochak · 2026-08-05T07:45:53Z

Maybe it's better to put this function into features.py ? I feel like 1 file with 1 function is too little, but it's up to you

okrochak · 2026-08-05T07:48:38Z

Again, a subjective thing, but maybe it's better to put ModelHubModelLoader into model-hub folder and import it in inference.py, to more clearly manage Model Hub code.

okrochak · 2026-08-05T08:16:20Z

I've tried to run this tutorial with a fresh install of itwinai from this branch modelhub-in-trainer. I don't think the .env file is uploaded, and when I try to use the .env file I have configured myself for the model hub, I still get an error even though .env file was parsed:

[ERROR]: API token not provided. Set it via: - --api-token option - HYPHA_TOKEN environment variable - HYPHA_TOKEN in .env file``` I will try to look into this deeper

r-sarma · 2026-08-25T12:33:32Z

@matbun Could you please review this?

matbun · 2026-08-25T19:54:20Z

@matbun Could you please review this?

I was waiting for you to iterate on Alex's comments, but I can review it in parallel. Will do by the end of the week!

matbun

Nice PR, it definitely improves a lot the integration with the AI model hub!

Important points (the order is random):

I would suggest moving the model_hub package from itwinai/torch/model_hub to one level up as itwinai/model_hub. The reason is that itwinai is naturally designed to support multiple frameworks, and model_hub should be considered at the same level as pytorch. Also considering that model hub does not strictly depend on torch. This also simplifies the import from itwinai.torch.model_hub to itwinai.model_hub
As already mentioned inline, I would remove the change in conftest.py and rebase on #505. This way we don't risk of having broken code while the CI still passes.
No test touches model_hub/*, ModelHubModelLoader, or the CLI changes. I would suggest adding some tests.
The Model Hub pull runs on every rank. TorchPredictor.execute calls self.model = self.model() without a rank guard, so all workers download the same checkpoint to the same CWD-relative tmp/modelhub_downloads/<model_id>/ path and race on the write. On HPC that path is usually the submit directory on a shared filesystem, so it's a cross-node corruption race, not just wasted bandwidth. The guard can't go inside ModelHubModelLoader since a ModelLoader is a bare callable with no strategy handle, but execute() method owns self.strategy: download under if self.strategy.is_main_worker: then self.strategy.barrier(), since every rank still needs the file before distribute_model()

Other points:

As discussed in a previous thread this could be a good opportunity of moving some code from upload_to_model_hub and _load_env_file to the new model_hub package to save some lines in cli.py . A similar reasoning is true for other functions we added in cli.py... If you agree I will open a separate issue for cleaning this up as well.

matbun · 2026-08-26T16:35:37Z

+os.environ.setdefault("MLFLOW_ALLOW_FILE_STORE", "true")
+


This change is making the CI pass, but it is hiding the root cause, which is being solved by #505
If the user doesn't remember to set this env var the code will break, which is something that the tests are not able to detect anymore.

I would suggest removing this, merging #505 first, and rebasing on main to bring in the fix

matbun · 2026-08-26T17:01:41Z

+    def __init__(self, config: dict):
+        self.config = config or {}
+        self.enabled = self.config.get("enabled", False)
+        self.final_checkpoint_name = self.config.get("final_checkpoint_name", "best_model")


final_checkpoint_name is configurable here, but trainer.py:1280 hardcodes the directory it
passes in:

best_ckpt_dir = Path(self.checkpoints_location) / "best_model"

So if a user sets final_checkpoint_name: my_ckpt, on_training_end hits the
ckpt_dir.name != self.final_checkpoint_name guard on line 29, returns early, and the whole
feature silently does nothing — no manifest, no upload, no message. The knob is only ever
correct at its default value.

Two ways out: drop the option and hardcode "best_model" in both places, or (better) have the
trainer ask the feature for the name it wants. The latter also removes the guard entirely:

# trainer.py if self.strategy.is_main_worker and self._model_hub.enabled: best_ckpt_dir = Path(self.checkpoints_location) / self._model_hub.final_checkpoint_name if best_ckpt_dir.exists(): self._model_hub.on_training_end(self, best_ckpt_dir)

Worth noting "best_model" is also hardcoded at trainer.py:1235 in the save_checkpoint
call, so a real fix probably wants a single constant both sites share.

matbun · 2026-08-26T17:03:29Z

+    def upload(self, model_dir: Path):
+        subprocess.run(
+            ["itwinai", "upload-model-to-hub", str(model_dir)],
+            check=False,


check=False means a non-zero exit from itwinai upload-model-to-hub never raises, right? in that case, this would make
_safe_upload's except Exception at feature.py:47 unreachable. so on a failed upload the
user never sees "Model Hub upload failed…" or the "You can re-upload later with…" hint, and
training reports success.

matbun · 2026-08-26T17:08:34Z

+        dst_dir = Path("tmp") / "modelhub_downloads" / self.model_id
+        ckpt_path = download_file(self.base_url, self.model_id, file_path, dst_dir)
+
+        checkpoint = torch.load(ckpt_path, weights_only=False)


I would suggest to set weights_only=True here, as it seems we are loading only the weights of the model. This has some important security implications considering we are unpicking a file we downloaded from the internet from a remote source which could have been compromised.

matbun · 2026-08-26T17:11:22Z

+        if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
+            model.load_state_dict(checkpoint["model_state_dict"], strict=False)
+        else:
+            model.load_state_dict(checkpoint, strict=False)


With strict=False, load_state_dict returns silently when no key matches at all. Combined
with pulling weights from a remote hub and requiring the user to supply model_class by hand,
this is the exact scenario where a mismatch is likely, and the result is a randomly
initialised model that runs inference and returns plausible-looking garbage... But if this is too strict please ignore this comment

matbun · 2026-09-11T12:10:41Z

Hi @r-sarma, I saw your last commits but I haven't started a review yet because I don't know if you are already done. Don't hesitate to re-request a review when ready. No rush, though

r-sarma · 2026-09-11T12:13:16Z

Hi @r-sarma, I saw your last commits but I haven't started a review yet because I don't know if you are already done. Don't hesitate to re-request a review when ready. No rush, though

Mostly done but I am adding some tests. I will re-request once its ready.

r-sarma added 5 commits May 8, 2026 10:17

adding model hub as feature

1d70595

fix linting errors

80663cb

editing build of manifest

77c8919

removing unused import

16344fa

adding model hub pulling API

60e8cc9

r-sarma self-assigned this Jul 21, 2026

r-sarma added the Feature extension label Jul 21, 2026

r-sarma added 6 commits July 29, 2026 14:42

reconfiguring pulling API download

a864dae

fix linting errors

6cb6b5f

adding tutorial, documentation and some refactoring

38bec8e

fix linting error and mlflow integration test

dc079c1

fix failing integration tests

5ad8856

Trigger Read the Docs rebuild

bc7f56c

r-sarma marked this pull request as ready for review August 5, 2026 07:24

r-sarma requested a review from matbun as a code owner August 5, 2026 07:24

r-sarma requested a review from okrochak August 5, 2026 07:24

okrochak reviewed Aug 5, 2026

View reviewed changes

matbun requested changes Aug 26, 2026

View reviewed changes

matbun and others added 10 commits September 9, 2026 14:53

Fix MLflow filesystem backend (#505)

248970d

Trigger Read the Docs rebuild

1905080

incorporating matbun comments

731cb03

incorporating matbun comments

cf66202

addressing matbun and okrochak PR comments

825cfba

fix paths

7d56d1d

changing weights_only argument for two loaders

8b856f1

rolling back changes to model loader

4fa8039

debugging integration test failure

a5fe612

moving model-hub related functions from src cli

e9bb520

		os.environ.setdefault("MLFLOW_ALLOW_FILE_STORE", "true")

Uh oh!

Conversation

r-sarma commented Jul 21, 2026

Uh oh!

okrochak left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

r-sarma commented Aug 25, 2026

Uh oh!

matbun commented Aug 25, 2026

Uh oh!

matbun left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

matbun commented Sep 11, 2026

Uh oh!

r-sarma commented Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants