def testUploadDefaultBoth(self):
other_local_dir = "/tmp/other"
other_cloud_dir = "memory:///other"
delete_at_uri(other_cloud_dir)
self._save_checkpoint_at(other_cloud_dir)
shutil.copytree(self.local_dir, other_local_dir)
# Case: Nothing is passed
checkpoint = TrialCheckpoint(local_path=self.local_dir,
cloud_path=self.cloud_dir)
path = checkpoint.upload()
self.assertEqual(self.cloud_dir, path)
# Case: Local dir is passed
checkpoint = TrialCheckpoint(local_path=self.local_dir,
cloud_path=self.cloud_dir)
path = checkpoint.upload(local_path=other_local_dir)
self.assertEqual(self.cloud_dir, path)
# Case: Both are passed
checkpoint = TrialCheckpoint(local_path=self.local_dir,
cloud_path=self.cloud_dir)
path = checkpoint.upload(local_path=other_local_dir,
cloud_path=other_cloud_dir)
self.assertEqual(other_cloud_dir, path)
def testUploadDefaultCloud(self):
other_cloud_dir = "memory:///other"
delete_at_uri(other_cloud_dir)
self._save_checkpoint_at(other_cloud_dir)
# Case: Nothing is passed
checkpoint = TrialCheckpoint(cloud_path=self.cloud_dir)
with self.assertRaisesRegex(RuntimeError, "No local path"):
checkpoint.upload()
# Case: Local dir is passed
checkpoint = TrialCheckpoint(cloud_path=self.cloud_dir)
path = checkpoint.upload(local_path=self.local_dir)
self.assertEqual(self.cloud_dir, path)
# Case: Cloud dir is passed
checkpoint = TrialCheckpoint(cloud_path=self.cloud_dir)
with self.assertRaisesRegex(RuntimeError, "No local path"):
checkpoint.upload(cloud_path=other_cloud_dir)
# Case: Both are passed
checkpoint = TrialCheckpoint(cloud_path=self.cloud_dir)
path = checkpoint.upload(local_path=self.local_dir,
cloud_path=other_cloud_dir)
self.assertEqual(other_cloud_dir, path)
def upload(
self,
cloud_path: Optional[str] = None,
local_path: Optional[str] = None,
clean_before: bool = False,
):
"""Upload checkpoint to cloud.
This will push the checkpoint directory from local storage
to ``cloud_path``.
If a ``cloud_path`` argument is provided and ``self.cloud_path``
is unset, it will be set to ``cloud_path``.
Args:
cloud_path: Cloud path to load checkpoint from.
Defaults to ``self.cloud_path``.
local_path: Local path to save checkpoint at.
Defaults to ``self.local_path``.
clean_before: If True, deletes potentially existing
cloud bucket before storing new data.
"""
local_path = local_path or self.local_path
if not local_path:
raise RuntimeError(
"Could not upload trial checkpoint: No local "
"path is set. Fix this by either passing a "
"`local_path` to your call to `upload()` or by "
"passing a `local_path` into the constructor."
)
cloud_path = cloud_path or self.cloud_path
if not cloud_path:
raise RuntimeError(
"Could not download trial checkpoint: No cloud "
"path is set. Fix this by either passing a "
"`cloud_path` to your call to `download()` or by "
"passing a `cloud_path` into the constructor. The latter "
"should automatically be done if you pass the correct "
"`tune.SyncConfig`."
)
if not self.cloud_path:
self.cloud_path = cloud_path
if clean_before:
logger.info(f"Clearing bucket contents before upload: {cloud_path}")
delete_at_uri(cloud_path)
# Actually upload
upload_to_uri(local_path, cloud_path)
return cloud_path
def test_fs_delete_at_uri(self):
"""Test that clear bucket utility works"""
checkpoint = self._prepare_fs_checkpoint()
# Convert into dict checkpoint
location = checkpoint.to_uri(self.cloud_uri)
delete_at_uri(location)
checkpoint = Checkpoint.from_uri(location)
with self.assertRaises(FileNotFoundError):
checkpoint.to_directory()
def test_fs_checkpoint_uri_pa(self):
"""Test conversion from fs to cloud checkpoint and back."""
checkpoint = self._prepare_fs_checkpoint()
# Clean up mock bucket
delete_at_uri(self.cloud_uri_pa)
_ensure_directory(self.cloud_uri_pa)
# Convert into dict checkpoint
location = checkpoint.to_uri(self.cloud_uri_pa)
self.assertIsInstance(location, str)
self.assertIn("mock://", location)
# Create from dict
checkpoint = Checkpoint.from_uri(location)
self.assertTrue(checkpoint._uri)
self._assert_fs_checkpoint(checkpoint)
def testSaveCloudTarget(self):
other_cloud_dir = "memory:///other"
delete_at_uri(other_cloud_dir)
self._save_checkpoint_at(other_cloud_dir)
# Case: No defaults
checkpoint = TrialCheckpoint()
with self.assertRaisesRegex(RuntimeError, "No existing local"):
checkpoint.save(self.cloud_dir)
# Case: Default local dir
# Write a checkpoint here as we assume existing local dir
with open(os.path.join(self.local_dir, "checkpoint.txt"), "wt") as f:
f.write("Checkpoint\n")
checkpoint = TrialCheckpoint(local_path=self.local_dir)
path = checkpoint.save(self.cloud_dir)
self.assertEqual(self.cloud_dir, path)
# Clean up checkpoint
os.remove(os.path.join(self.local_dir, "checkpoint.txt"))
# Case: Default cloud dir, copy to other cloud
checkpoint = TrialCheckpoint(cloud_path=self.cloud_dir)
path = checkpoint.save(other_cloud_dir)
self.assertEqual(other_cloud_dir, path)
# Case: Default both, copy to other cloud
checkpoint = TrialCheckpoint(local_path=self.local_dir,
cloud_path=self.cloud_dir)
path = checkpoint.save(other_cloud_dir)
self.assertEqual(other_cloud_dir, path)
def delete_external_checkpoint(checkpoint_uri: str):
delete_at_uri(checkpoint_uri)
def _save_checkpoint_at(self, target):
delete_at_uri(target)
upload_to_uri(local_path=self.local_dir, uri=target)
def tearDown(self) -> None:
shutil.rmtree(self.local_dir)
delete_at_uri(self.cloud_dir)
