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

nshtrainer/_checkpoint/metadata.py

@@ -111,11 +111,11 @@ def _write_checkpoint_metadata(
     try:
         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}")
+        log.exception(f"Failed to write metadata to {metadata_path}")
         return None
 
-    log.debug(f"Checkpoint metadata written to {checkpoint_path}")
-    return checkpoint_path
+    log.debug(f"Checkpoint metadata written to {metadata_path}")
+    return metadata_path
 
 
 def _remove_checkpoint_metadata(checkpoint_path: Path):

nshtrainer/_hf_hub.py

@@ -1,3 +1,4 @@
+import io
 import logging
 import os
 from pathlib import Path
@@ -299,7 +300,22 @@ def _save_checkpoint_files(
     # Resolve the repository name
     repo_name = _repo_name(api, root_config)
 
+    # Let's read all the files to memory right now,
+    # in case they get used/removed by other processes.
+    # Read all the files to memory
+    file_contents: list[bytes | None] = []
     for p in paths:
+        try:
+            with open(p, "rb") as f:
+                file_contents.append(f.read())
+        except IOError as e:
+            log.warning(f"Failed to read checkpoint file {p}: {str(e)}")
+            file_contents.append(None)
+
+    for p, contents in zip(paths, file_contents):
+        if contents is None:
+            continue
+
         try:
             relative_path = p.relative_to(checkpoint_dir)
         except ValueError:
@@ -314,7 +330,7 @@ def _save_checkpoint_files(
         # Upload the checkpoint file to the repository
         try:
             api.upload_file(
-                path_or_fileobj=str(p.resolve().absolute()),
+                path_or_fileobj=io.BytesIO(contents),
                 path_in_repo=str(path_in_repo),
                 repo_id=repo_name,
                 repo_type="model",

nshtrainer/trainer/trainer.py

@@ -426,7 +426,7 @@ class Trainer(LightningTrainer):
             metadata_path = _write_checkpoint_metadata(self, lm, filepath)
 
         # If HF Hub is enabled, then we upload
-        if hparams.trainer.hf_hub:
+        if hparams.trainer.hf_hub and self.is_global_zero:
             from .._hf_hub import _save_checkpoint_files
 
             files = [f for f in (filepath, metadata_path) if f is not None]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.18.0
+Version: 0.18.2
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

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

@@ -1,9 +1,9 @@
 nshtrainer/__init__.py,sha256=39loiLLXbaGiozEsAn8mPHopxaPsek8JsgR9DD2gxtY,583
 nshtrainer/_checkpoint/loader.py,sha256=myFObRsPdb8jBncMK73vjr5FDJIfKhF86Ec_kSjXtwg,13837
-nshtrainer/_checkpoint/metadata.py,sha256=_9dBLJSCgi3H98-HJLgwVr8U7yHxbQA5VB9ZYMYjFj0,5181
+nshtrainer/_checkpoint/metadata.py,sha256=p5e7dhVPpOGrXeuesq_7Y_RHi5lguzDAR_UXtMJXzWU,5175
 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/_hf_hub.py,sha256=Py9_8ADvMCFPaJzeE7bxm8Mgs3mEMkyWJ4pDEccTGt8,11230
 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
@@ -77,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=xJBl8C-9SVT1ppmxTVwT1PIN8vZmE1erpKtKlsX2-8Y,17479
+nshtrainer/trainer/trainer.py,sha256=TTtVkgSB_ekgDlHg24d58Vzddtkpp6ZHOTVprXdXMH0,17503
 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
@@ -85,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.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,,
+nshtrainer-0.18.2.dist-info/METADATA,sha256=vev96DaxCnqJOAvvGrGOJ37OpWNFLrCdtGPN-kpnvO4,935
+nshtrainer-0.18.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.18.2.dist-info/RECORD,,