Is your feature request related to a problem? Please describe. Currently, the save_state() and load_state() methods only work with file paths, requiring the optimizer state to be written to the local file system. This makes it impossible to save state directly to cloud storage solutions (e.g., AWS S3, Google Cloud Storage, Azure Blob Storage) without creating temporary local files. In distributed computing environments or serverless architectures where local file system access is limited or ephemeral, this limitation requires additional workarounds and introduces unnecessary I/O overhead.
Describe the solution you'd like Extract the state serialization logic into separate methods that work with dictionaries/JSON objects directly:
_state_to_dict() — Returns the optimizer state as a Python dictionary without writing to a file_load_state_dict(state: dict) — Loads optimizer state from a dictionary without reading from a file
The existing save_state() and load_state() methods should be refactored to use these new internal methods, maintaining backward compatibility.
This enables direct cloud storage integration:
# Save to S3 without temporary files
state_dict = optimizer._state_to_dict()
s3_client.put_object(Bucket="bucket", Key="state.json", Body=json.dumps(state_dict))
# Load from S3 directly
obj = s3_client.get_object(Bucket="bucket", Key="state.json")
optimizer._load_state_dict(json.load(obj["Body"]))
Yes, I can provide this feature.
Is your feature request related to a problem? Please describe. Currently, the
save_state()andload_state()methods only work with file paths, requiring the optimizer state to be written to the local file system. This makes it impossible to save state directly to cloud storage solutions (e.g., AWS S3, Google Cloud Storage, Azure Blob Storage) without creating temporary local files. In distributed computing environments or serverless architectures where local file system access is limited or ephemeral, this limitation requires additional workarounds and introduces unnecessary I/O overhead.Describe the solution you'd like Extract the state serialization logic into separate methods that work with dictionaries/JSON objects directly:
_state_to_dict()— Returns the optimizer state as a Python dictionary without writing to a file_load_state_dict(state: dict)— Loads optimizer state from a dictionary without reading from a fileThe existing
save_state()andload_state()methods should be refactored to use these new internal methods, maintaining backward compatibility.This enables direct cloud storage integration:
Yes, I can provide this feature.