Method save

internlm/utils/storage_manager.py:483–501

(self, save_path: str, saved_obj: Any, *args, async_upload=None, **kwargs)

        return meta.client.get_fns(*unpack_meta(meta))

    def save(self, save_path: str, saved_obj: Any, *args, async_upload=None, **kwargs):
        meta = self._get_client(path=save_path)

        if async_upload is None:
            async_upload = self.async_mode
        if async_upload:
            assert (
                self.tmp_local_folder
            ), "StorageManager is not setted tmp_local_folder, so async save cannot be performed."
            tmp_step_file = meta.local_nvme_path
            self._to_be_del_files.append(tmp_step_file)
            with open(tmp_step_file, "wb") as f:
                torch.save(saved_obj, f, pickle_protocol=pickle.HIGHEST_PROTOCOL)
            self.async_executor(meta.async_upload_fn, *unpack_meta(meta))
            os.chmod(tmp_step_file, stat.S_IRWXU | stat.S_IRWXG | stat.S_IRWXO)
            self.async_task_peeding = True
        else:
            meta.client.sync_upload_fileobj(*unpack_meta(meta), *args, saved_obj=saved_obj, **kwargs)
            self.upload_count += 1

    def load(self, load_path: str, *args, **kwargs) -> Any:
        self.wait()