Update to Model Hub integration

docs/api/itwinai.torch.modules.rst

@@ -47,6 +47,60 @@ loggers
    :member-order: bysource
 
 
+model_hub.feature
+++++++++++++++++++
+.. automodule:: itwinai.torch.model_hub.feature
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :member-order: bysource
+
+
+model_hub.download
++++++++++++++++++++
+.. automodule:: itwinai.torch.model_hub.download
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :member-order: bysource
+
+
+model_hub.manifest
++++++++++++++++++++
+.. automodule:: itwinai.torch.model_hub.manifest
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :member-order: bysource
+
+
+model_hub.utils
+++++++++++++++++
+.. automodule:: itwinai.torch.model_hub.utils
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :member-order: bysource
+
+
+model_hub.backends.base
+++++++++++++++++++++++++
+.. automodule:: itwinai.torch.model_hub.backends.base
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :member-order: bysource
+
+
+model_hub.backends.itwinai_hub
++++++++++++++++++++++++++++++++
+.. automodule:: itwinai.torch.model_hub.backends.itwinai_hub
+   :members:
+   :undoc-members:
+   :show-inheritance:
+   :member-order: bysource
+
+
 models.mnist
 ++++++++++++
 .. automodule:: itwinai.torch.models.mnist

docs/how-it-works/model-hub/explain_model_hub.rst

@@ -0,0 +1,103 @@
+Accessing models from the RI-SCALE Model Hub
+============================================
+
+**Author(s)**: Rakesh Sarma (FZJ)
+
+Once a ML model has been trained, it is often needed to be shared with collaborators, or to
+publish it in open repositories. itwinai integrates with the `RI-SCALE Model Hub
+<https://modelhub.riscale.eu>`_ to support users with this functionality which allows pushing
+a trained checkpoint to the Model Hub, and pulling a checkpoint to run inference on it.
+
+Pushing a model
+---------------
+
+Pushing is handled automatically by :class:`~itwinai.torch.trainer.TorchTrainer` whenever
+Model Hub support is enabled in its configuration. On every checkpoint save
+(:meth:`~itwinai.torch.trainer.TorchTrainer.save_checkpoint`), the checkpoint directory --
+containing ``model.pt`` (the model's raw ``state_dict``), ``state.pt`` (optimizer/scheduler/
+epoch state), and ``config.yaml`` -- is handed to
+:class:`~itwinai.torch.model_hub.feature.ModelHubFeature`, which:
+
+1. Writes a ``manifest.yaml`` into the checkpoint directory via
+   :func:`~itwinai.torch.model_hub.manifest.write_manifest`, merging user-supplied fields
+   with sensible defaults. At minimum, ``id`` and ``name`` must be provided.
+2. Uploads the checkpoint directory using the configured backend at the end of all epochs.
+   Backends implement :class:`~itwinai.torch.model_hub.backends.base.BaseBackend` and are
+   selected by name via :func:`~itwinai.torch.model_hub.backends.get_backend`. Currently the
+   only backend is :class:`~itwinai.torch.model_hub.backends.itwinai_hub.AIModelHubBackend`.
+   The abstraction is to enable future backends (e.g. HuggingFace).
+
+The timing of the upload is controlled by a ``mode`` setting:
+
+- ``online``: upload immediately, regardless of connectivity.
+- ``auto``: upload if internet is available; otherwise print the checkpoint's local
+  location and skip the upload.
+- ``deferred``: never upload automatically; the checkpoint is left ready to be pushed
+  manually later.
+
+.. admonition:: Example Model Hub push configuration
+
+    .. code-block:: yaml
+
+        model_hub:
+          enabled: true
+          backend: ai-model-hub
+          mode: online
+          manifest:
+            id: checkpoint-example
+            name: My Model
+            published: true
+
+    The final `published: true` ensures that the pushed model is readily visible to all users
+    on the AI Model Hub.
+
+Pulling a model
+---------------
+
+Pulling is handled by :class:`~itwinai.torch.inference.ModelHubModelLoader`, an
+implementation of :class:`~itwinai.serialization.ModelLoader`. Like any other
+``ModelLoader``, it can be used wherever a model loader is expected -- most commonly as the
+``model`` argument of :class:`~itwinai.torch.inference.TorchPredictor`.
+
+Unlike pushing, the Model Hub's file API has no endpoint to download a whole model folder
+at once: files are retrieved one at a time, by exact path
+(``GET /artifacts/{model_id}/files/{file_path}``). To spare users from needing to know that
+exact path, ``ModelHubModelLoader`` supports two modes:
+
+  - If ``file_path`` is provided explicitly, that file is downloaded directly.
+  - If ``file_path`` is omitted, itwinai lists the model's files
+    (:func:`~itwinai.torch.model_hub.download.list_files`) and locates
+    ``root/<checkpoint_dir_name>/model.pt`` automatically
+    (:func:`~itwinai.torch.model_hub.download.discover_weights_file`), matching the layout
+    produced by :meth:`~itwinai.torch.trainer.TorchTrainer.save_checkpoint`.
+    ``discover_weights_file`` only looks for a top-level ``root/`` entry; it does not
+    inspect or otherwise handle any other top-level entries the Hub may contain.
+
+.. admonition:: Example Model Hub pull configuration
+
+    .. code-block:: yaml
+
+        predictor:
+          _target_: itwinai.torch.inference.TorchPredictor
+          config: {}
+          model:
+            _target_: itwinai.torch.inference.ModelHubModelLoader
+            model_id: checkpoint-example
+            model_class: my_module.MyModel
+
+.. important::
+   Model Hub checkpoints store a raw ``state_dict`` -- just tensors, with no architecture
+   information -- following the same convention used by
+   :meth:`~itwinai.torch.trainer.TorchTrainer.save_checkpoint`. This means ``model_class``
+   is **always required** when pulling: it must be the exact :class:`~torch.nn.Module`
+   subclass used at training time. If the original training script is not available, the
+   downloaded ``state_dict``'s keys and tensor shapes can be inspected directly to
+   reconstruct a matching class by hand.
+
+Connectivity
+------------
+
+Both pushing (in ``auto`` mode) and pulling rely on the same connectivity check,
+:func:`~itwinai.torch.model_hub.utils.has_internet_connection`. Pulling always requires
+internet access -- unlike pushing, there is no offline or deferred mode for pulling, since
+there is no local fallback artifact to use in its place.

docs/index.rst

@@ -32,6 +32,7 @@
    how-it-works/workflows/explain_workflows
    how-it-works/hpo/explain-hpo
    how-it-works/scalability-report/scalability_report
+   how-it-works/model-hub/explain_model_hub
 
 .. toctree::
    :maxdepth: 2
@@ -145,6 +146,7 @@ Quick Start
 - :ref:`Hyper-parameter Optimization <hpo-tutorials>`
 - :ref:`ML Workflows <ml-workflows-tutorials>`
 - :ref:`Code Profiling and Optimization <profiling-tutorials>`
+- :ref:`Model Hub <model-hub-tutorials>`
 
 📚 Use Cases & 🧩 Plugins
 ==========================

docs/tutorials/model-hub/model_hub_tutorial.rst

@@ -0,0 +1,33 @@
+Pushing and pulling models with the RI-SCALE Model Hub
+======================================================
+.. include:: ../../../tutorials/model-hub/torch-tutorial-model-hub/README.md
+   :parser: myst_parser.sphinx_
+   :start-line: 4
+
+Run the training pipeline (trains the model and, as per ``model_hub`` in ``config.yaml``,
+pushes the best checkpoint to the Model Hub):
+
+.. code-block:: bash
+
+   itwinai exec-pipeline +pipe-key training_pipeline
+
+Then run the inference pipeline (pulls that same checkpoint and runs inference on it):
+
+.. code-block:: bash
+
+   itwinai exec-pipeline +pipe-key inference_pipeline
+
+config.yaml
++++++++++++
+.. literalinclude:: ../../../tutorials/model-hub/torch-tutorial-model-hub/config.yaml
+   :language: yaml
+
+data.py
++++++++
+.. literalinclude:: ../../../tutorials/model-hub/torch-tutorial-model-hub/data.py
+   :language: python
+
+synthetic_data.py
++++++++++++++++++
+.. literalinclude:: ../../../tutorials/model-hub/torch-tutorial-model-hub/synthetic_data.py
+   :language: python

docs/tutorials/tutorials.rst

@@ -48,16 +48,16 @@ Here you can find a collection of tutorials for various complexity ML workflows.
 
 .. toctree::
    :maxdepth: 1
-   
+
    workflows/01-pipeline-introduction/tutorial_0_basic_workflow
    workflows/02-pipeline-configuration/tutorial_1_intermediate_workflow
    workflows/03-dag-workflows/tutorial_2_advanced_workflow
    workflows/04_itwinai_argparser
- 
+
 
 .. _hpo-tutorials:
 
-Hyperparameter Optimization 
+Hyperparameter Optimization
 ===========================
 
 This tutorial provides an overview of Hyperparameter Optimization (HPO) workflows.
@@ -81,3 +81,16 @@ Here you can find our tutorials on how to do profiling with **itwinai**:
    profiling/profiling-overview
    profiling/py-spy-profiling
    profiling/py-spy-lattice-qcd-example
+
+
+.. _model-hub-tutorials:
+
+Model Hub
+=========
+
+Here you can find our tutorial on pushing and pulling models with the RI-SCALE Model Hub.
+
+.. toctree::
+   :maxdepth: 1
+
+   model-hub/model_hub_tutorial

src/itwinai/cli.py

@@ -551,6 +551,7 @@ def _load_slurm_builder_config(
 
     return validated_config, (root_config if expect_nested else None)
 
+
 @app.command()
 def run(
     config: Annotated[
@@ -976,7 +977,7 @@ def upload_model_to_hub(
         str | None,
         typer.Option(
             "--env-file",
-            help="Path to .env file containing MODEL_HUB_URL and MODEL_HUB_API_TOKEN.",
+            help="Path to .env file containing HYPHA_SERVER_URL and HYPHA_TOKEN.",
         ),
     ] = None,
     upload_script: Annotated[
@@ -1006,29 +1007,24 @@ def _check_internet_connection(timeout: float = 3.0) -> bool:
         import socket
 
         try:
-            socket.create_connection(("8.8.8.8", 53), timeout=timeout)
+            socket.create_connection(("1.1.1.1", 443), timeout=timeout)
             return True
         except OSError:
             return False
 
     model_path = Path(model_dir).resolve()
 
     # Validate if model directory exists and is a directory
-    if not model_path.exists():
-        cli_logger.error(f"Model directory '{model_path}' does not exist!")
-        raise typer.Exit(code=1)
-
-    if not model_path.is_dir():
-        cli_logger.error(f"'{model_path}' is not a directory!")
+    if not model_path.exists() or not model_path.is_dir():
+        cli_logger.error(
+            f"Model directory '{model_path}' does not exist or is not a directory."
+        )
         raise typer.Exit(code=1)
 
     # Check if the file manifest.yaml exists
     manifest_file = model_path / "manifest.yaml"
     if not manifest_file.exists():
-        cli_logger.error(
-            f"No manifest.yaml found in '{model_path}'. "
-            "The model directory must contain a manifest.yaml file with the model id."
-        )
+        cli_logger.error(f"No manifest.yaml found in '{model_path}'. ")
         raise typer.Exit(code=1)
 
     # Load environment variables from .env file if specified and if file exists
@@ -1121,16 +1117,30 @@ def _check_internet_connection(timeout: float = 3.0) -> bool:
     cli_logger.info(f"Uploading model from '{model_path}' to {final_hub_url}")
 
     try:
-        # Call the upload script as subprocess
-        # The original usage is: python upload_model.py model_example1
-        result = subprocess.run(
-            [sys.executable, str(upload_script_path), str(model_path)],
-            env=env_vars,
-            capture_output=True,
-            text=True,
-            cwd=str(upload_script_path.parent),  # Run from script directory
-            check=False,
-        )
+        # Build the "root/<ckpt_dir_name>" layout that discover_weights_file expects,
+        # via a symlink
+        scratch_dir = Path(tempfile.mkdtemp())
+        root_dir = scratch_dir / "root"
+        root_dir.mkdir(parents=True, exist_ok=True)
+        symlink_path = root_dir / model_path.name
+        if not symlink_path.exists():
+            symlink_path.symlink_to(model_path, target_is_directory=True)
+
+        relative_upload_arg = f"root/{model_path.name}"
+
+        try:
+            # Call the upload script as subprocess
+            # The original usage is: python upload_model.py model_example1
+            result = subprocess.run(
+                [sys.executable, str(upload_script_path), relative_upload_arg],
+                env=env_vars,
+                capture_output=True,
+                text=True,
+                cwd=str(scratch_dir),
+                check=False,
+            )
+        finally:
+            shutil.rmtree(scratch_dir, ignore_errors=True)
 
         # Print stdout (even if there's an error, this may be useful for debugging)
         if result.stdout:
@@ -1171,12 +1181,13 @@ def _load_env_file(env_path: Path, env_dict: dict):
                 key = key.strip()
                 value = value.strip()
                 # Remove quotes if present
-                if (
-                    (value.startswith('"') and value.endswith('"')) or
-                    (value.startswith("'") and value.endswith("'"))
+                if (value.startswith('"') and value.endswith('"')) or (
+                    value.startswith("'") and value.endswith("'")
                 ):
                     value = value[1:-1]
 
+                env_dict[key] = value
+
 
 if __name__ == "__main__":
     app()