chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,636 @@
|
||||
import logging
|
||||
import os
|
||||
import pickle
|
||||
import re
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import pyarrow.fs
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray import train, tune
|
||||
from ray._common.test_utils import simulate_s3_bucket
|
||||
from ray.air._internal.uri_utils import URI
|
||||
from ray.air.constants import EXPR_RESULT_FILE
|
||||
from ray.train._checkpoint import Checkpoint
|
||||
from ray.train._internal.storage import (
|
||||
StorageContext,
|
||||
_delete_fs_path,
|
||||
_download_from_fs_path,
|
||||
)
|
||||
from ray.train.base_trainer import TrainingFailedError
|
||||
from ray.train.data_parallel_trainer import DataParallelTrainer
|
||||
from ray.tune.trainable.trainable import _DICT_CHECKPOINT_FILE_NAME
|
||||
|
||||
|
||||
class TestConstants:
|
||||
NUM_ITERATIONS = 6 # == num_checkpoints == num_artifacts
|
||||
NUM_TRIALS = 2
|
||||
NUM_WORKERS = 3
|
||||
|
||||
SCORE_KEY = "score"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def mock_s3_bucket_uri():
|
||||
port = 5002
|
||||
region = "us-west-2"
|
||||
with simulate_s3_bucket(port=port, region=region) as s3_uri:
|
||||
import boto3
|
||||
|
||||
s3 = boto3.client(
|
||||
"s3", region_name=region, endpoint_url=f"http://localhost:{port}"
|
||||
)
|
||||
# Bucket name will be autogenerated/unique per test
|
||||
bucket_name = URI(s3_uri).name
|
||||
s3.create_bucket(
|
||||
Bucket=bucket_name,
|
||||
CreateBucketConfiguration={"LocationConstraint": region},
|
||||
)
|
||||
# Disable server HTTP request logging
|
||||
logging.getLogger("werkzeug").setLevel(logging.WARNING)
|
||||
yield URI(s3_uri)
|
||||
logging.getLogger("werkzeug").setLevel(logging.INFO)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def dummy_context_manager(*args, **kwargs):
|
||||
yield "dummy value"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope="module")
|
||||
def ray_start_4_cpus():
|
||||
ray.init(num_cpus=4)
|
||||
yield
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def _create_mock_custom_fs(custom_fs_root_dir: Path) -> pyarrow.fs.FileSystem:
|
||||
from fsspec.implementations.dirfs import DirFileSystem
|
||||
from fsspec.implementations.local import LocalFileSystem
|
||||
|
||||
custom_fs_root_dir.mkdir(parents=True, exist_ok=True)
|
||||
storage_filesystem = pyarrow.fs.PyFileSystem(
|
||||
pyarrow.fs.FSSpecHandler(
|
||||
DirFileSystem(path=str(custom_fs_root_dir), fs=LocalFileSystem())
|
||||
)
|
||||
)
|
||||
return storage_filesystem
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _resolve_storage_type(
|
||||
storage_path_type: str, tmp_path: Path
|
||||
) -> Tuple[str, Optional[pyarrow.fs.FileSystem]]:
|
||||
storage_path, storage_filesystem = None, None
|
||||
|
||||
context_manager = (
|
||||
mock_s3_bucket_uri if storage_path_type == "cloud" else dummy_context_manager
|
||||
)
|
||||
|
||||
with context_manager() as cloud_storage_path:
|
||||
if storage_path_type == "nfs":
|
||||
storage_path = str(tmp_path / "fake_nfs")
|
||||
elif storage_path_type == "cloud":
|
||||
storage_path = str(cloud_storage_path)
|
||||
elif storage_path_type == "custom_fs":
|
||||
storage_path = "mock_bucket"
|
||||
storage_filesystem = _create_mock_custom_fs(tmp_path / "custom_fs")
|
||||
|
||||
yield storage_path, storage_filesystem
|
||||
|
||||
|
||||
def _get_local_inspect_dir(
|
||||
root_local_path: Path,
|
||||
storage_path: str,
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem],
|
||||
storage_local_path: Path = None,
|
||||
) -> Tuple[Path, str]:
|
||||
"""Downloads the storage path -> local dir for inspecting contents.
|
||||
|
||||
Args:
|
||||
root_local_path: Local directory to use as the inspect root.
|
||||
storage_path: The storage path or URI to download from.
|
||||
storage_filesystem: Optional custom filesystem to use.
|
||||
storage_local_path: Local path that ``storage_path`` mirrors on disk
|
||||
when no remote storage is configured.
|
||||
|
||||
Returns:
|
||||
Tuple: (local_inspect_dir, storage_fs_path), where storage_fs_path
|
||||
is the path to the storage path on the filesystem (e.g., prefix stripped).
|
||||
This is used to check the correctness of paths returned from `Result`'s,
|
||||
since URIs are hard to do comparisons with.
|
||||
"""
|
||||
local_inspect_dir = root_local_path / "inspect"
|
||||
if storage_path:
|
||||
if storage_filesystem:
|
||||
fs, storage_fs_path = storage_filesystem, storage_path
|
||||
else:
|
||||
fs, storage_fs_path = pyarrow.fs.FileSystem.from_uri(storage_path)
|
||||
_download_from_fs_path(
|
||||
fs=fs, fs_path=storage_fs_path, local_path=str(local_inspect_dir)
|
||||
)
|
||||
else:
|
||||
fs, storage_fs_path = pyarrow.fs.LocalFileSystem(), str(storage_local_path)
|
||||
local_inspect_dir = storage_local_path
|
||||
|
||||
return local_inspect_dir, storage_fs_path
|
||||
|
||||
|
||||
def _get_checkpoint_index(checkpoint_dir_name: str) -> int:
|
||||
"""Gets the checkpoint index from the checkpoint directory name."""
|
||||
return int(checkpoint_dir_name.split("_")[-1])
|
||||
|
||||
|
||||
def _create_checkpoint_shard_filename(rank_str: str) -> str:
|
||||
return f"checkpoint_shard-rank={rank_str}.pkl"
|
||||
|
||||
|
||||
def _get_checkpoint_shard_rank(checkpoint_shard_filename: str) -> int:
|
||||
"""Get the checkpoint shard rank from the filename."""
|
||||
pattern = _create_checkpoint_shard_filename(r"(\d+)")
|
||||
match = re.search(pattern, checkpoint_shard_filename)
|
||||
assert match
|
||||
return int(match.group(1))
|
||||
|
||||
|
||||
def train_fn(config):
|
||||
in_trainer = config.get("in_trainer", False)
|
||||
if in_trainer:
|
||||
from ray.train._internal.session import _TrainSession, get_session
|
||||
|
||||
train_session = get_session()
|
||||
|
||||
assert isinstance(train_session, _TrainSession)
|
||||
assert train_session.storage
|
||||
assert train_session.storage.checkpoint_fs_path
|
||||
|
||||
# Check that the working dir for each worker is the shared trial dir.
|
||||
assert (
|
||||
Path.cwd() == Path(train_session.storage.trial_working_directory).resolve()
|
||||
)
|
||||
|
||||
start = 0
|
||||
|
||||
checkpoint = train.get_checkpoint()
|
||||
if checkpoint:
|
||||
custom_restore_fn = config.get("custom_restore_fn")
|
||||
if custom_restore_fn:
|
||||
state = custom_restore_fn(checkpoint)
|
||||
else:
|
||||
with checkpoint.as_directory() as checkpoint_dir:
|
||||
with open(os.path.join(checkpoint_dir, "checkpoint.pkl"), "rb") as f:
|
||||
state = pickle.load(f)
|
||||
print("Loaded back state from checkpoint:", state)
|
||||
start = state["iter"] + 1
|
||||
|
||||
for i in range(start, config.get("num_iterations", 5)):
|
||||
time.sleep(config.get("time_per_iter", 0.25))
|
||||
|
||||
metrics = {"iter": i, TestConstants.SCORE_KEY: i}
|
||||
|
||||
# Save an artifact in the local trial dir.
|
||||
rank = train.get_context().get_world_rank()
|
||||
artifact_file_name = (
|
||||
f"artifact-rank={rank}-iter={i}.txt"
|
||||
if in_trainer
|
||||
else f"artifact-iter={i}.txt"
|
||||
)
|
||||
with open(artifact_file_name, "w") as f:
|
||||
f.write(f"{i}")
|
||||
|
||||
if in_trainer and train.get_context().get_world_rank() in config.get(
|
||||
"no_checkpoint_ranks", []
|
||||
):
|
||||
train.report(metrics)
|
||||
else:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
with open(os.path.join(temp_dir, "checkpoint.pkl"), "wb") as f:
|
||||
pickle.dump({"iter": i}, f)
|
||||
|
||||
if in_trainer:
|
||||
checkpoint_file_name = _create_checkpoint_shard_filename(str(rank))
|
||||
with open(os.path.join(temp_dir, checkpoint_file_name), "wb") as f:
|
||||
pickle.dump({"iter": i}, f)
|
||||
|
||||
with config.get("custom_save_fn", dummy_context_manager)(temp_dir):
|
||||
train.report(
|
||||
metrics, checkpoint=Checkpoint.from_directory(temp_dir)
|
||||
)
|
||||
# `train.report` should not have deleted this!
|
||||
assert os.path.exists(temp_dir)
|
||||
|
||||
if i in config.get("fail_iters", []):
|
||||
raise RuntimeError(f"Failing on iter={i}!!")
|
||||
|
||||
|
||||
class ClassTrainable(tune.Trainable):
|
||||
"""Implement (almost) the same thing as `train_fn` but as a class."""
|
||||
|
||||
def setup(self, config):
|
||||
# Save some markers in the trial dir.
|
||||
tmp_path = config.get("tmp_path")
|
||||
self.fail_markers = {
|
||||
i: tmp_path / f"fail_marker_{self.trial_id}_iter={i}"
|
||||
for i in config.get("fail_iters", [])
|
||||
}
|
||||
setup_marker = tmp_path / f"setup_marker_{self.trial_id}"
|
||||
if not setup_marker.exists():
|
||||
for marker in self.fail_markers.values():
|
||||
marker.touch()
|
||||
setup_marker.touch()
|
||||
|
||||
self.save_as_dict = config.get("save_checkpoint_as_dict", False)
|
||||
|
||||
def step(self) -> dict:
|
||||
if self.iteration in self.fail_markers:
|
||||
marker = self.fail_markers[self.iteration]
|
||||
if marker.exists():
|
||||
marker.unlink()
|
||||
raise RuntimeError(f"Failing on iter={self.iteration}")
|
||||
|
||||
# Save an artifact in the local trial dir.
|
||||
artifact_file_name = f"artifact-iter={self.iteration}.txt"
|
||||
with open(artifact_file_name, "w") as f:
|
||||
f.write(f"{self.iteration}")
|
||||
|
||||
return {
|
||||
"score": 1,
|
||||
"done": self.iteration >= self.config.get("num_iterations") - 1,
|
||||
"should_checkpoint": True,
|
||||
}
|
||||
|
||||
def save_checkpoint(self, temp_checkpoint_dir) -> str:
|
||||
if self.save_as_dict:
|
||||
return {"dummy": "data"}
|
||||
(Path(temp_checkpoint_dir) / "checkpoint.pkl").write_text("dummy")
|
||||
return temp_checkpoint_dir
|
||||
|
||||
def load_checkpoint(self, checkpoint_dict_or_path):
|
||||
print("Loading state from:", checkpoint_dict_or_path)
|
||||
print("At iteration =", self.iteration)
|
||||
if self.save_as_dict:
|
||||
assert checkpoint_dict_or_path == {"dummy": "data"}
|
||||
else:
|
||||
assert (
|
||||
Path(checkpoint_dict_or_path) / "checkpoint.pkl"
|
||||
).read_text() == "dummy"
|
||||
|
||||
|
||||
def _resume_from_checkpoint(
|
||||
checkpoint: Checkpoint,
|
||||
expected_state: dict,
|
||||
storage_path: Optional[str] = None,
|
||||
storage_filesystem: Optional[pyarrow.fs.FileSystem] = None,
|
||||
):
|
||||
print(f"\nStarting run with `resume_from_checkpoint`: {checkpoint}\n")
|
||||
|
||||
def assert_fn(config):
|
||||
checkpoint_to_check = train.get_checkpoint()
|
||||
with checkpoint_to_check.as_directory() as checkpoint_dir:
|
||||
with open(os.path.join(checkpoint_dir, "checkpoint.pkl"), "rb") as f:
|
||||
state = pickle.load(f)
|
||||
|
||||
print("Loaded state from `resume_from_checkpoint`:", state)
|
||||
print("Expected state:", expected_state)
|
||||
assert state == expected_state, (state, expected_state)
|
||||
|
||||
dummy_ckpt = tempfile.mkdtemp()
|
||||
with open(os.path.join(dummy_ckpt, "dummy.txt"), "w") as f:
|
||||
f.write("data")
|
||||
train.report({"dummy": 1}, checkpoint=Checkpoint.from_directory(dummy_ckpt))
|
||||
|
||||
trainer = DataParallelTrainer(
|
||||
assert_fn,
|
||||
scaling_config=train.ScalingConfig(num_workers=2),
|
||||
run_config=train.RunConfig(
|
||||
name="test_resume_from_checkpoint",
|
||||
storage_path=storage_path,
|
||||
storage_filesystem=storage_filesystem,
|
||||
),
|
||||
resume_from_checkpoint=checkpoint,
|
||||
)
|
||||
result = trainer.fit()
|
||||
|
||||
# Make sure that the checkpoint indexing starts from scratch.
|
||||
assert Path(
|
||||
result.checkpoint.path
|
||||
).name == StorageContext._make_checkpoint_dir_name(0)
|
||||
|
||||
# Clean up this run's experiment directory immediately after.
|
||||
_delete_fs_path(result.filesystem, Path(result.path).parent.as_posix())
|
||||
|
||||
|
||||
def _assert_storage_contents(
|
||||
local_inspect_dir: Path,
|
||||
exp_name: str,
|
||||
checkpoint_config: train.CheckpointConfig,
|
||||
trainable_name: str,
|
||||
test_trainer: bool,
|
||||
no_checkpoint_ranks: List[int] = None,
|
||||
constants: type = TestConstants,
|
||||
):
|
||||
no_checkpoint_ranks = no_checkpoint_ranks or []
|
||||
|
||||
# Second, inspect the contents of the storage path
|
||||
storage_path_ls = list(local_inspect_dir.glob("*"))
|
||||
assert len(storage_path_ls) == 1 # Only expect 1 experiment dir
|
||||
exp_dir = storage_path_ls[0]
|
||||
assert exp_dir.name == exp_name
|
||||
|
||||
# Files synced by the driver
|
||||
assert len(list(exp_dir.glob("tuner.pkl"))) == 1
|
||||
if test_trainer:
|
||||
assert len(list(exp_dir.glob("trainer.pkl"))) == 1
|
||||
# 2 copies of these files:
|
||||
# 1 for the initial run, and 1 for the manually restored run.
|
||||
assert len(list(exp_dir.glob("basic-variant-state-*"))) == 2
|
||||
assert len(list(exp_dir.glob("experiment_state-*"))) == 2
|
||||
|
||||
# Files synced by the worker
|
||||
assert (
|
||||
len(list(exp_dir.glob(f"{trainable_name}*"))) == 1
|
||||
if test_trainer
|
||||
else constants.NUM_TRIALS
|
||||
)
|
||||
for trial_dir in exp_dir.glob(f"{trainable_name}*"):
|
||||
# If set, expect num_to_keep. Otherwise, expect to see all of them.
|
||||
expected_num_checkpoints = (
|
||||
checkpoint_config.num_to_keep or constants.NUM_ITERATIONS
|
||||
)
|
||||
|
||||
assert len(list(trial_dir.glob("checkpoint_*"))) == expected_num_checkpoints
|
||||
checkpoint_idxs = sorted(
|
||||
[
|
||||
_get_checkpoint_index(checkpoint_dir.name)
|
||||
for checkpoint_dir in trial_dir.glob("checkpoint_*")
|
||||
]
|
||||
)
|
||||
# Ex: If num_to_keep=2 out of 6 total checkpoints,
|
||||
# expect checkpoint_004 and checkpoint_005.
|
||||
assert checkpoint_idxs == list(
|
||||
range(
|
||||
constants.NUM_ITERATIONS - expected_num_checkpoints,
|
||||
constants.NUM_ITERATIONS,
|
||||
)
|
||||
)
|
||||
|
||||
for checkpoint_dir in trial_dir.glob("checkpoint_*"):
|
||||
# 1 shared checkpoint.pkl file, written by the trainable / all workers.
|
||||
assert (
|
||||
len(list(checkpoint_dir.glob("checkpoint.pkl"))) == 1
|
||||
# NOTE: Dict checkpoint is only for the ClassTrainable.
|
||||
or len(list(checkpoint_dir.glob(_DICT_CHECKPOINT_FILE_NAME))) == 1
|
||||
)
|
||||
if test_trainer:
|
||||
# 1 checkpoint shard per worker.
|
||||
# Unless the worker did not report a checkpoint (no_checkpoint_ranks).
|
||||
assert {
|
||||
_get_checkpoint_shard_rank(checkpoint_shard.name)
|
||||
for checkpoint_shard in checkpoint_dir.glob(
|
||||
"checkpoint_shard-*.pkl"
|
||||
)
|
||||
} == {
|
||||
i
|
||||
for i in range(constants.NUM_WORKERS)
|
||||
if i not in no_checkpoint_ranks
|
||||
}
|
||||
|
||||
if test_trainer:
|
||||
expected_num_artifacts = constants.NUM_ITERATIONS * constants.NUM_WORKERS
|
||||
else:
|
||||
expected_num_artifacts = constants.NUM_ITERATIONS
|
||||
assert len(list(trial_dir.glob("artifact-*"))) == expected_num_artifacts
|
||||
|
||||
# NOTE: This result file is synced by the driver.
|
||||
assert len(list(trial_dir.glob(EXPR_RESULT_FILE))) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("trainable", [train_fn, ClassTrainable])
|
||||
@pytest.mark.parametrize("storage_path_type", ["nfs", "cloud", "custom_fs"])
|
||||
@pytest.mark.parametrize(
|
||||
"checkpoint_config",
|
||||
[train.CheckpointConfig(), train.CheckpointConfig(num_to_keep=2)],
|
||||
)
|
||||
def test_tuner(
|
||||
tmp_path,
|
||||
trainable,
|
||||
storage_path_type,
|
||||
checkpoint_config: train.CheckpointConfig,
|
||||
):
|
||||
"""End-to-end test that the new persistence mode works with the Tuner API.
|
||||
This test covers many `storage_path_type` options:
|
||||
- storage_path=None --> save locally to the default local path (e.g., ~/ray_results)
|
||||
- storage_path="nfs" --> save locally to a fake NFS path
|
||||
- storage_path="cloud" --> save to a mock S3 bucket
|
||||
- storage_path="custom_fs" --> save to a custom pyarrow filesystem
|
||||
- The custom fs is a local filesystem that appends a path prefix to every path.
|
||||
|
||||
This is the expected output at the storage path:
|
||||
|
||||
{storage_path}/{exp_name}
|
||||
├── tuner.pkl <- Driver artifacts (global experiment state)
|
||||
├── basic-variant-state.json
|
||||
├── experiment_state.json
|
||||
├── train_fn_a2b9e_00000_0_...
|
||||
│ ├── artifact-iter=0.txt <- Trial artifacts
|
||||
│ ├── ...
|
||||
│ ├── checkpoint_000000 <- Trial checkpoints
|
||||
│ │ └── checkpoint.pkl
|
||||
│ ├── ...
|
||||
│ ├── events.out.tfevents... <- Driver artifacts (trial results)
|
||||
│ ├── params.json
|
||||
│ ├── params.pkl
|
||||
│ ├── progress.csv
|
||||
│ └── result.json
|
||||
└── train_fn_a2b9e_00001_1_...
|
||||
└── ... <- Same as above
|
||||
"""
|
||||
exp_name = f"tuner_persistence_test-{uuid.uuid4().hex}"
|
||||
|
||||
with _resolve_storage_type(storage_path_type, tmp_path) as (
|
||||
storage_path,
|
||||
storage_filesystem,
|
||||
):
|
||||
run_config = train.RunConfig(
|
||||
storage_path=storage_path,
|
||||
storage_filesystem=storage_filesystem,
|
||||
name=exp_name,
|
||||
verbose=0,
|
||||
failure_config=train.FailureConfig(max_failures=1),
|
||||
checkpoint_config=checkpoint_config,
|
||||
sync_config=train.SyncConfig(sync_artifacts=True),
|
||||
)
|
||||
tuner = tune.Tuner(
|
||||
trainable,
|
||||
param_space={
|
||||
"num_iterations": TestConstants.NUM_ITERATIONS,
|
||||
"fail_iters": [2, 4],
|
||||
# NOTE: This param is only used in the ClassTrainable.
|
||||
"save_checkpoint_as_dict": tune.grid_search([True, False]),
|
||||
"tmp_path": tmp_path,
|
||||
},
|
||||
run_config=run_config,
|
||||
# 2 samples (from the grid search). Run 1 at at time to test actor reuse
|
||||
tune_config=tune.TuneConfig(num_samples=1, max_concurrent_trials=1),
|
||||
)
|
||||
result_grid = tuner.fit()
|
||||
assert result_grid.errors
|
||||
|
||||
restored_tuner = tune.Tuner.restore(
|
||||
path=str(URI(run_config.storage_path) / exp_name),
|
||||
trainable=trainable,
|
||||
storage_filesystem=storage_filesystem,
|
||||
resume_errored=True,
|
||||
)
|
||||
result_grid = restored_tuner.fit()
|
||||
assert not result_grid.errors
|
||||
|
||||
local_inspect_dir, storage_fs_path = _get_local_inspect_dir(
|
||||
root_local_path=tmp_path,
|
||||
storage_path=run_config.storage_path,
|
||||
storage_filesystem=storage_filesystem,
|
||||
)
|
||||
|
||||
# First, check that the ResultGrid returns the correct paths.
|
||||
print(result_grid)
|
||||
experiment_fs_path = result_grid.experiment_path
|
||||
assert isinstance(result_grid.filesystem, pyarrow.fs.FileSystem), result_grid
|
||||
assert experiment_fs_path == os.path.join(storage_fs_path, exp_name)
|
||||
assert len(result_grid) == TestConstants.NUM_TRIALS
|
||||
for result in result_grid:
|
||||
trial_fs_path = result.path
|
||||
assert isinstance(result.filesystem, pyarrow.fs.FileSystem), result
|
||||
assert trial_fs_path.startswith(experiment_fs_path)
|
||||
for checkpoint, _ in result.best_checkpoints:
|
||||
assert checkpoint.path.startswith(trial_fs_path)
|
||||
|
||||
# Next, inspect the storage path contents.
|
||||
_assert_storage_contents(
|
||||
local_inspect_dir,
|
||||
exp_name,
|
||||
checkpoint_config,
|
||||
trainable_name=trainable.__name__,
|
||||
test_trainer=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("storage_path_type", ["nfs", "cloud", "custom_fs"])
|
||||
@pytest.mark.parametrize(
|
||||
"checkpoint_config",
|
||||
[
|
||||
train.CheckpointConfig(),
|
||||
train.CheckpointConfig(
|
||||
num_to_keep=1,
|
||||
checkpoint_score_attribute=TestConstants.SCORE_KEY,
|
||||
checkpoint_score_order="max",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_trainer(
|
||||
tmp_path, storage_path_type, checkpoint_config: train.CheckpointConfig
|
||||
):
|
||||
"""Same end-to-end test as `test_tuner`, but also includes a
|
||||
`DataParallelTrainer(resume_from_checkpoint)` test at the end.
|
||||
|
||||
{storage_path}/{exp_name}
|
||||
├── experiment_state-2023-07-28_10-00-38.json <- Initial exp state
|
||||
├── basic-variant-state-2023-07-28_10-00-38.json
|
||||
├── experiment_state-2023-07-28_10-01-38.json <- Restored exp state
|
||||
├── basic-variant-state-2023-07-28_10-01-38.json
|
||||
├── trainer.pkl
|
||||
├── tuner.pkl
|
||||
└── DataParallelTrainer_46367_00000_0_...
|
||||
├── events.out.tfevents...
|
||||
├── params.json
|
||||
├── params.pkl
|
||||
├── progress.csv
|
||||
├── result.json
|
||||
├── checkpoint_000000
|
||||
│ ├── checkpoint.pkl <- Shared checkpoint file
|
||||
│ ├── checkpoint_shard-rank=0.pkl <- Worker checkpoint shards
|
||||
│ └── checkpoint_shard-rank=1.pkl
|
||||
├── ...
|
||||
├── artifact-rank=0-iter=0.txt <- Worker artifacts
|
||||
├── artifact-rank=1-iter=0.txt
|
||||
├── ...
|
||||
├── artifact-rank=0-iter=1.txt
|
||||
├── artifact-rank=1-iter=1.txt
|
||||
└── ...
|
||||
"""
|
||||
exp_name = f"trainer_persistence_test-{uuid.uuid4().hex}"
|
||||
no_checkpoint_ranks = [0]
|
||||
|
||||
with _resolve_storage_type(storage_path_type, tmp_path) as (
|
||||
storage_path,
|
||||
storage_filesystem,
|
||||
):
|
||||
run_config = train.RunConfig(
|
||||
storage_path=storage_path,
|
||||
storage_filesystem=storage_filesystem,
|
||||
name=exp_name,
|
||||
verbose=0,
|
||||
checkpoint_config=checkpoint_config,
|
||||
failure_config=train.FailureConfig(max_failures=1),
|
||||
sync_config=train.SyncConfig(sync_artifacts=True),
|
||||
)
|
||||
trainer = DataParallelTrainer(
|
||||
train_fn,
|
||||
train_loop_config={
|
||||
"in_trainer": True,
|
||||
"num_iterations": TestConstants.NUM_ITERATIONS,
|
||||
"fail_iters": [2, 4],
|
||||
# Test that global rank 0 is not required to checkpoint.
|
||||
"no_checkpoint_ranks": no_checkpoint_ranks,
|
||||
},
|
||||
scaling_config=train.ScalingConfig(num_workers=TestConstants.NUM_WORKERS),
|
||||
run_config=run_config,
|
||||
)
|
||||
print("\nStarting initial run.\n")
|
||||
with pytest.raises(TrainingFailedError):
|
||||
result = trainer.fit()
|
||||
|
||||
print("\nStarting manually restored run.\n")
|
||||
restored_trainer = DataParallelTrainer.restore(
|
||||
path=str(URI(run_config.storage_path) / exp_name),
|
||||
storage_filesystem=storage_filesystem,
|
||||
)
|
||||
result = restored_trainer.fit()
|
||||
|
||||
_resume_from_checkpoint(
|
||||
result.checkpoint,
|
||||
expected_state={"iter": TestConstants.NUM_ITERATIONS - 1},
|
||||
)
|
||||
|
||||
local_inspect_dir, storage_fs_path = _get_local_inspect_dir(
|
||||
root_local_path=tmp_path,
|
||||
storage_path=run_config.storage_path,
|
||||
storage_filesystem=storage_filesystem,
|
||||
)
|
||||
|
||||
# First, inspect that the result object returns the correct paths.
|
||||
print(result)
|
||||
trial_fs_path = result.path
|
||||
assert trial_fs_path.startswith(storage_fs_path)
|
||||
for checkpoint, _ in result.best_checkpoints:
|
||||
assert checkpoint.path.startswith(trial_fs_path)
|
||||
|
||||
_assert_storage_contents(
|
||||
local_inspect_dir,
|
||||
exp_name,
|
||||
checkpoint_config,
|
||||
trainable_name="DataParallelTrainer",
|
||||
test_trainer=True,
|
||||
no_checkpoint_ranks=no_checkpoint_ranks,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Reference in New Issue
Block a user