Function function_trainable

python/ray/tune/utils/release_test_util.py:76–103

(config)

def function_trainable(config):
    num_iters = int(config["num_iters"])
    sleep_time = config["sleep_time"]
    score = config["score"]

    checkpoint_iters = config["checkpoint_iters"]
    checkpoint_size_b = config["checkpoint_size_b"]
    checkpoint_num_items = checkpoint_size_b // 8  # np.float64
    checkpoint_num_files = config["checkpoint_num_files"]

    for i in range(num_iters):
        metrics = {"score": i + score}
        if (
            checkpoint_iters >= 0
            and checkpoint_size_b > 0
            and i % checkpoint_iters == 0
        ):
            with tempfile.TemporaryDirectory() as tmpdir:
                for i in range(checkpoint_num_files):
                    checkpoint_file = os.path.join(tmpdir, f"bogus_{i}.ckpt")
                    checkpoint_data = np.random.uniform(0, 1, size=checkpoint_num_items)
                    with open(checkpoint_file, "wb") as fp:
                        pickle.dump(checkpoint_data, fp)
                tune.report(metrics, checkpoint=Checkpoint.from_directory(tmpdir))
        else:
            tune.report(metrics)

        time.sleep(sleep_time)


def timed_tune_run(