nshtrainer 0.17.1__py3-none-any.whl → 0.18.0__py3-none-any.whl

nshtrainer/_checkpoint/metadata.py

@@ -112,8 +112,10 @@ def _write_checkpoint_metadata(
         metadata_path.write_text(metadata.model_dump_json(indent=4), encoding="utf-8")
     except Exception:
         log.exception(f"Failed to write metadata to {checkpoint_path}")
-    else:
-        log.debug(f"Checkpoint metadata written to {checkpoint_path}")
+        return None
+
+    log.debug(f"Checkpoint metadata written to {checkpoint_path}")
+    return checkpoint_path
 
 
 def _remove_checkpoint_metadata(checkpoint_path: Path):

nshtrainer/_hf_hub.py

@@ -0,0 +1,347 @@
+import logging
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import nshconfig as C
+from lightning.pytorch import Callback
+from lightning.pytorch.trainer import Trainer
+from nshrunner._env import SNAPSHOT_DIR
+from typing_extensions import override
+
+from .callbacks.base import (
+    CallbackConfigBase,
+    CallbackMetadataConfig,
+    CallbackWithMetadata,
+)
+
+if TYPE_CHECKING:
+    from huggingface_hub import HfApi  # noqa: F401
+
+    from .model.base import BaseConfig
+    from .trainer.trainer import Trainer
+log = logging.getLogger(__name__)
+
+
+class HuggingFaceHubAutoCreateConfig(C.Config):
+    enabled: bool = True
+    """Enable automatic repository creation on the Hugging Face Hub."""
+
+    private: bool = True
+    """Whether to create the repository as private."""
+
+    namespace: str | None = None
+    """The namespace to create the repository in. If `None`, the repository will be created in the user's namespace."""
+
+    def __bool__(self):
+        return self.enabled
+
+
+class HuggingFaceHubConfig(CallbackConfigBase):
+    """Configuration options for Hugging Face Hub integration."""
+
+    enabled: bool = False
+    """Enable Hugging Face Hub integration."""
+
+    token: str | None = None
+    """Hugging Face Hub API token. If `None`, the token will be read from the current environment.
+    This needs to either be set using `huggingface-cli login` or by setting the `HUGGINGFACE_TOKEN`
+    environment variable."""
+
+    auto_create: HuggingFaceHubAutoCreateConfig = HuggingFaceHubAutoCreateConfig()
+    """Automatic repository creation configuration options."""
+
+    save_config: bool = True
+    """Whether to save the model configuration to the Hugging Face Hub."""
+
+    save_checkpoints: bool = True
+    """Whether to save checkpoints to the Hugging Face Hub."""
+
+    save_code: bool = True
+    """Whether to save code to the Hugging Face Hub.
+    This is only supported if `nshsnap` is installed and snapshotting is enabled."""
+
+    save_in_background: bool = True
+    """Whether to save to the Hugging Face Hub in the background.
+    This corresponds to setting `run_as_future=True` in the HFApi upload methods."""
+
+    def enable_(self):
+        self.enabled = True
+        return self
+
+    def disable_(self):
+        self.enabled = False
+        return self
+
+    def __bool__(self):
+        return self.enabled
+
+    @override
+    def create_callbacks(self, root_config):
+        yield CallbackWithMetadata(
+            HFHubCallback(self),
+            CallbackMetadataConfig(ignore_if_exists=True),
+        )
+
+
+def _api(token: str | None = None):
+    # Make sure that `huggingface_hub` is installed
+    try:
+        import huggingface_hub  # noqa: F401
+    except ImportError:
+        log.exception(
+            "Could not import `huggingface_hub`. Please install it using `pip install huggingface_hub`."
+        )
+        return None
+
+    # Create and authenticate the API instance
+    try:
+        api = huggingface_hub.HfApi(token=token)
+
+        # Verify authentication
+        api.whoami()
+    except Exception as e:
+        log.exception(
+            f"Authentication failed for Hugging Face Hub: {str(e)}. "
+            "Please make sure you are logged in using `huggingface-cli login`, "
+            "by setting the HUGGING_FACE_HUB_TOKEN environment variable, "
+            "or by providing a valid token in the configuration."
+        )
+        return None
+
+    return api
+
+
+def _enabled_and_valid(
+    trainer: "Trainer",
+    config: HuggingFaceHubConfig,
+    *,
+    rank_zero_only: bool,
+):
+    # Make sure this is enabled and the config is valid
+    if not config:
+        return None
+
+    # If `rank_zero_only` and this is not rank 0, stop here.
+    if rank_zero_only and not trainer.is_global_zero:
+        return
+
+    # Make sure that `huggingface_hub` is installed
+    try:
+        import huggingface_hub  # noqa: F401
+    except ImportError:
+        log.exception(
+            "Could not import `huggingface_hub`. Please install it using `pip install huggingface_hub`."
+        )
+        return None
+
+    # Create and authenticate the API instance
+    if (api := getattr(trainer, "_hf_hub_api", None)) is None:
+        api = _api(config.token)
+        setattr(trainer, "_hf_hub_api", api)
+    return cast(huggingface_hub.HfApi, api)
+
+
+def _repo_name(api: "HfApi", root_config: "BaseConfig"):
+    username = None
+    if (ac := root_config.trainer.hf_hub.auto_create) and ac.namespace:
+        username = ac.namespace
+    elif (username := api.whoami().get("name", None)) is None:
+        raise ValueError("Could not get username from Hugging Face Hub.")
+
+    return f"{username}/{root_config.project}-{root_config.run_name}-{root_config.id}"
+
+
+def _init(*, trainer: "Trainer", root_config: "BaseConfig"):
+    config = root_config.trainer.hf_hub
+    if (
+        api := _enabled_and_valid(
+            trainer,
+            config,
+            rank_zero_only=True,
+        )
+    ) is None or not config.auto_create:
+        return
+
+    from huggingface_hub.utils import RepositoryNotFoundError
+
+    # Resolve the repository name
+    repo_name = _repo_name(api, root_config)
+
+    # Create the repository, if it doesn't exist
+    try:
+        # Check if the repository exists
+        api.repo_info(repo_id=repo_name, repo_type="model")
+        log.info(f"Repository '{repo_name}' already exists.")
+    except RepositoryNotFoundError:
+        # Repository doesn't exist, so create it
+        try:
+            api.create_repo(
+                repo_id=repo_name,
+                repo_type="model",
+                private=config.auto_create.private,
+                exist_ok=True,
+            )
+            log.info(f"Created new repository '{repo_name}'.")
+        except Exception as e:
+            log.exception(f"Failed to create repository '{repo_name}': {str(e)}")
+    except Exception as e:
+        log.exception(f"Error checking repository '{repo_name}': {str(e)}")
+
+    # Upload the config
+    _save_config(root_config, trainer=trainer)
+
+    # Upload the code
+    _save_code(repo_name, config=config, trainer=trainer)
+
+
+def _save_code(
+    repo_name: str,
+    *,
+    config: HuggingFaceHubConfig,
+    trainer: "Trainer",
+):
+    if (
+        api := _enabled_and_valid(
+            trainer,
+            config,
+            rank_zero_only=True,
+        )
+    ) is None or not config.save_code:
+        return
+
+    # If a snapshot has been taken (which can be detected using the SNAPSHOT_DIR env),
+    # then upload all contents within the snapshot directory to the repository.
+    snapshot_dir = os.environ.get(SNAPSHOT_DIR)
+    if not snapshot_dir:
+        log.info("No snapshot directory found. Skipping upload.")
+        return
+
+    snapshot_path = Path(snapshot_dir)
+    if not snapshot_path.exists() or not snapshot_path.is_dir():
+        log.warning(
+            f"Snapshot directory '{snapshot_dir}' does not exist or is not a directory."
+        )
+        return
+
+    try:
+        api.upload_folder(
+            folder_path=str(snapshot_path),
+            repo_id=repo_name,
+            repo_type="model",
+            path_in_repo="code",  # Prefix with "code" folder
+            run_as_future=cast(Any, config.save_in_background),
+        )
+        log.info(
+            f"Uploaded snapshot contents to repository '{repo_name}' under 'code' folder."
+        )
+    except Exception as e:
+        log.exception(
+            f"Failed to upload snapshot contents to repository '{repo_name}' under 'code' folder: {str(e)}"
+        )
+
+
+def _save_config(
+    root_config: "BaseConfig",
+    *,
+    trainer: "Trainer",
+):
+    config = root_config.trainer.hf_hub
+    if (
+        api := _enabled_and_valid(
+            trainer,
+            config,
+            rank_zero_only=True,
+        )
+    ) is None or not config.save_config:
+        return
+
+    # Convert the root config to a JSON string
+    # NOTE: This is a utf-8 string.
+    config_json = root_config.model_dump_json(indent=4)
+
+    # Resolve the repository name
+    repo_name = _repo_name(api, root_config)
+
+    # Upload the config file to the repository
+    try:
+        api.upload_file(
+            path_or_fileobj=config_json.encode("utf-8"),
+            path_in_repo="config.json",
+            repo_id=repo_name,
+            repo_type="model",
+            run_as_future=cast(Any, config.save_in_background),
+        )
+        log.info(f"Uploaded config.json to repository '{repo_name}'.")
+    except Exception as e:
+        log.exception(
+            f"Failed to upload config.json to repository '{repo_name}': {str(e)}"
+        )
+
+
+def _save_checkpoint_files(
+    trainer: "Trainer",
+    paths: list[Path],
+    *,
+    root_config: "BaseConfig",
+):
+    config = root_config.trainer.hf_hub
+    if (
+        api := _enabled_and_valid(trainer, config, rank_zero_only=True)
+    ) is None or not config.save_checkpoints:
+        return
+
+    # Resolve the checkpoint directory
+    checkpoint_dir = root_config.directory.resolve_subdirectory(
+        root_config.id, "checkpoint"
+    )
+
+    # Resolve the repository name
+    repo_name = _repo_name(api, root_config)
+
+    for p in paths:
+        try:
+            relative_path = p.relative_to(checkpoint_dir)
+        except ValueError:
+            log.warning(
+                f"Checkpoint path {p} is not within the checkpoint directory {checkpoint_dir}."
+            )
+            continue
+
+        # Prefix the path in repo with "checkpoints"
+        path_in_repo = Path("checkpoints") / relative_path
+
+        # Upload the checkpoint file to the repository
+        try:
+            api.upload_file(
+                path_or_fileobj=str(p.resolve().absolute()),
+                path_in_repo=str(path_in_repo),
+                repo_id=repo_name,
+                repo_type="model",
+                run_as_future=cast(Any, config.save_in_background),
+            )
+            log.info(
+                f"Uploaded checkpoint file {relative_path} to repository '{repo_name}'."
+            )
+        except Exception as e:
+            log.exception(
+                f"Failed to upload checkpoint file {relative_path} to repository '{repo_name}': {str(e)}"
+            )
+
+    log.info(f"Completed uploading checkpoint files to repository '{repo_name}'.")
+
+
+class HFHubCallback(Callback):
+    def __init__(self, config: HuggingFaceHubConfig):
+        super().__init__()
+        self.config = config
+
+    @override
+    def setup(self, trainer, pl_module, stage):
+        root_config = cast("BaseConfig", pl_module.hparams)
+        _init(trainer=trainer, root_config=root_config)
+
+    @override
+    def teardown(self, trainer, pl_module, stage):
+        if hasattr(trainer, "_hf_hub_api"):
+            delattr(trainer, "_hf_hub_api")

nshtrainer/model/__init__.py

@@ -10,6 +10,7 @@ from .config import CheckpointSavingConfig as CheckpointSavingConfig
 from .config import DirectoryConfig as DirectoryConfig
 from .config import EarlyStoppingConfig as EarlyStoppingConfig
 from .config import GradientClippingConfig as GradientClippingConfig
+from .config import HuggingFaceHubConfig as HuggingFaceHubConfig
 from .config import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
 from .config import LoggingConfig as LoggingConfig
 from .config import MetricConfig as MetricConfig

nshtrainer/model/base.py

@@ -192,6 +192,7 @@ class LightningModuleBase(  # pyright: ignore[reportIncompatibleMethodOverride]
             hparams = self.config_cls().model_validate(hparams)
         hparams.environment = EnvironmentConfig.from_current_environment(hparams, self)
         hparams = self.pre_init_update_hparams(hparams)
+
         super().__init__(hparams)
 
         self.save_hyperparameters(hparams)

nshtrainer/model/config.py

@@ -33,6 +33,7 @@ from lightning.pytorch.strategies.strategy import Strategy
 from typing_extensions import Self, TypedDict, TypeVar, override
 
 from .._checkpoint.loader import CheckpointLoadingConfig
+from .._hf_hub import HuggingFaceHubConfig
 from ..callbacks import (
     BestCheckpointCallbackConfig,
     CallbackConfig,
@@ -819,6 +820,9 @@ class TrainerConfig(C.Config):
     checkpoint_saving: CheckpointSavingConfig = CheckpointSavingConfig()
     """Checkpoint saving configuration options."""
 
+    hf_hub: HuggingFaceHubConfig = HuggingFaceHubConfig()
+    """Hugging Face Hub configuration options."""
+
     logging: LoggingConfig = LoggingConfig()
     """Logging/experiment tracking (e.g., WandB) configuration options."""
 
@@ -1213,4 +1217,5 @@ class BaseConfig(C.Config):
         yield self.trainer.checkpoint_saving
         yield self.trainer.logging
         yield self.trainer.optimizer
+        yield self.trainer.hf_hub
         yield from self.trainer.callbacks

nshtrainer/trainer/trainer.py

@@ -420,8 +420,16 @@ class Trainer(LightningTrainer):
         # Save the checkpoint metadata
         lm = self._base_module
         hparams = cast(BaseConfig, lm.hparams)
+        metadata_path = None
         if hparams.trainer.save_checkpoint_metadata and self.is_global_zero:
             # Generate the metadata and write to disk
-            _write_checkpoint_metadata(self, lm, filepath)
+            metadata_path = _write_checkpoint_metadata(self, lm, filepath)
+
+        # If HF Hub is enabled, then we upload
+        if hparams.trainer.hf_hub:
+            from .._hf_hub import _save_checkpoint_files
+
+            files = [f for f in (filepath, metadata_path) if f is not None]
+            _save_checkpoint_files(self, files, root_config=hparams)
 
         return ret_val

{nshtrainer-0.17.1.dist-info → nshtrainer-0.18.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.17.1
+Version: 0.18.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -11,6 +11,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: extra
 Requires-Dist: GitPython ; extra == "extra"
+Requires-Dist: huggingface-hub ; extra == "extra"
 Requires-Dist: lightning
 Requires-Dist: nshconfig
 Requires-Dist: nshrunner

{nshtrainer-0.17.1.dist-info → nshtrainer-0.18.0.dist-info}/RECORD

@@ -1,8 +1,9 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_checkpoint/loader.py,sha256=myFObRsPdb8jBncMK73vjr5FDJIfKhF86Ec_kSjXtwg,13837
-nshtrainer/_checkpoint/metadata.py,sha256=FGVYqqHp5rCETcPfaoSZmGIPapE4kdYJCKSutTRERQI,5147
+nshtrainer/_checkpoint/metadata.py,sha256=_9dBLJSCgi3H98-HJLgwVr8U7yHxbQA5VB9ZYMYjFj0,5181
 nshtrainer/_checkpoint/saver.py,sha256=DkbCH0YeOJ71m32vAARiQdGBf0hvwwdoAV8LOFGy-0Y,1428
 nshtrainer/_experimental/__init__.py,sha256=pEXPyI184UuDHvfh4p9Kg9nQZQZI41e4_HvNd4BK-yg,81
+nshtrainer/_hf_hub.py,sha256=b1Na0-SyOM5xlJCH8cqjk0ggEVCPMI_z770c32JIQRY,10701
 nshtrainer/callbacks/__init__.py,sha256=4qocBDzQbLLhhbIEfvbA3SQB_Dy9ZJH7keMwPay-ZS8,2359
 nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHmPDtzdC4ctgXyRudkRJqH4m4,23184
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
@@ -54,9 +55,9 @@ nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=mn6cyizyI_stkXtg6zxIEGF9b
 nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=h76oTHYpMxauV_l6lviya5DW-WKArwxxf7ZQizhmbCw,2782
 nshtrainer/metrics/__init__.py,sha256=ObLIELGguIEcUpRsUkqh1ltrvZii6vglTpJGrPvoy00,50
 nshtrainer/metrics/_config.py,sha256=jgRBfDAQLFTW7AiUY7CRtdfts6CR6keeuqm0FFMWCzQ,1288
-nshtrainer/model/__init__.py,sha256=BmqSbf6v6oyeilti4iEn_Tyrr1kRmcFcJekTb8NeglI,1315
-nshtrainer/model/base.py,sha256=AXRfEsFAT0Ln7zjYVPU5NgtHS_c8FZM-M4pyLamO7OA,17516
-nshtrainer/model/config.py,sha256=D6Y-Y7GoMrpo7A2dmIqJsqc4X2IHwyl9OEHxO4uOc0g,42918
+nshtrainer/model/__init__.py,sha256=VyRziPT3YilP6xjLi_StsSqtlvn7N4LOMzgukRsOnF8,1380
+nshtrainer/model/base.py,sha256=oQVolDk81acy4OlckwQEBHuX2gCaVSYiIA0JaDIfhQ4,17517
+nshtrainer/model/config.py,sha256=147uV7IukvuYE4G_ZuQNxVjnlog1BdCrAVbcj_sx9Vs,43104
 nshtrainer/model/modules/callback.py,sha256=K0-cyEtBcQhI7Q2e-AGTE8T-GghUPY9DYmneU6ULV6g,6401
 nshtrainer/model/modules/debug.py,sha256=Yy7XEdPou9BkCsD5hJchwJGmCVGrfUru5g9VjPM4uAw,1120
 nshtrainer/model/modules/distributed.py,sha256=ABpR9d-3uBS_fivfy_WYW-dExW6vp5BPaoPQnOudHng,1725
@@ -76,7 +77,7 @@ nshtrainer/trainer/__init__.py,sha256=P2rmr8oBVTHk-HJHYPcUwWqDEArMbPR4_rPpATbWK3
 nshtrainer/trainer/_runtime_callback.py,sha256=sd2cUdRJG-UCdQr9ruZvEYpNGNF1t2W2fuxwwVlQD9E,4164
 nshtrainer/trainer/checkpoint_connector.py,sha256=F2tkHogbMAa5U7335sm77sZBkjEDa5v46XbJCH9Mg6c,2167
 nshtrainer/trainer/signal_connector.py,sha256=2EzkVktlasl8PgWAKNLDZRUMY__gRlDy1HdinAU-tfU,10740
-nshtrainer/trainer/trainer.py,sha256=jIqiNrq1I0f5pQP7lHshtgjCAYfpoWPoqwS74LHU9iM,17148
+nshtrainer/trainer/trainer.py,sha256=xJBl8C-9SVT1ppmxTVwT1PIN8vZmE1erpKtKlsX2-8Y,17479
 nshtrainer/util/_environment_info.py,sha256=gIdq9TJgzGCdcVzZxjHcwYasJ_HmEGVHbvE-KJVVtWs,24187
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
 nshtrainer/util/environment.py,sha256=AeW_kLl-N70wmb6L_JLz1wRj0kA70xs6RCmc9iUqczE,4159
@@ -84,6 +85,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.17.1.dist-info/METADATA,sha256=yFOvBuWFb0naLc6p82HVodctRPi1AAzWIb3wwa5Vq-I,885
-nshtrainer-0.17.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.17.1.dist-info/RECORD,,
+nshtrainer-0.18.0.dist-info/METADATA,sha256=uKcju9SCdP6M3h-GjX0OOcpd52_cNThmUPmMYUpBIk4,935
+nshtrainer-0.18.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.18.0.dist-info/RECORD,,