How to use the ray.tune.run function in ray

checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
        torch.save(self.model.state_dict(), checkpoint_path)
        return checkpoint_path

    def _restore(self, checkpoint_path):
        self.model.load_state_dict(torch.load(checkpoint_path))


# __trainable_example_end__
# yapf: enable

if __name__ == "__main__":
    args = parser.parse_args()
    ray.init(address=args.ray_address)
    sched = ASHAScheduler(metric="mean_accuracy")
    analysis = tune.run(
        TrainMNIST,
        scheduler=sched,
        stop={
            "mean_accuracy": 0.95,
            "training_iteration": 3 if args.smoke_test else 20,
        },
        resources_per_trial={
            "cpu": 3,
            "gpu": int(args.use_gpu)
        },
        num_samples=1 if args.smoke_test else 20,
        checkpoint_at_end=True,
        checkpoint_freq=3,
        config={
            "args": args,
            "lr": tune.uniform(0.001, 0.1),

def run(model_args, objective, teacher_model, result_dir, nmaxepochs, grace_period):
    experiment = distillation_experiment()
    try:
        with open('../config/redis_address', 'r') as f:
            address = f.read().strip()
            ray.init(redis_address=address)
    except:
        ray.init()
    ahb = AsyncHyperBandScheduler(metric='mean_loss', mode='min', grace_period=grace_period,
                                  max_t=nmaxepochs)
                                  #  reduction_factor=2, brackets=3, max_t=nmaxepochs)
    trials = ray.tune.run(experiment, scheduler=ahb, raise_on_failed_trial=False,
                          queue_trials=True, reuse_actors=True).trials
    trials = [trial for trial in trials if trial.last_result is not None]
    loss = [trial.last_result.get('mean_loss', float('inf')) for trial in trials]

    return teacher_model, model_args, objective, min(loss)

def run(model, result_dir, nmaxepochs, grace_period):
    experiment = distillation_experiment()
    try:
        with open('../config/redis_address', 'r') as f:
            address = f.read().strip()
            ray.init(redis_address=address)
    except:
        ray.init()
    ahb = AsyncHyperBandScheduler(reward_attr='inverse_loss', grace_period=grace_period, reduction_factor=2, brackets=3, max_t=nmaxepochs)
    trials = ray.tune.run(experiment, scheduler=ahb, raise_on_failed_trial=False, queue_trials=True)
    trials = [trial for trial in trials if trial.last_result is not None]
    loss = [trial.last_result.get('mean_loss', float('inf')) for trial in trials]

    checkpoint_path = Path(result_dir) / experiment.name
    checkpoint_path.mkdir(parents=True, exist_ok=True)
    checkpoint_path /= 'trial.pkl'
    with checkpoint_path.open('wb') as f:
        pickle.dump(trials, f)

    ex.add_artifact(str(checkpoint_path))
    return min(loss)

else:
        ray.init()

    pbt = PopulationBasedTraining(
        time_attr="training_iteration",
        metric="mean_accuracy",
        mode="max",
        perturbation_interval=20,
        hyperparam_mutations={
            # distribution for resampling
            "lr": lambda: random.uniform(0.0001, 0.02),
            # allow perturbations within this set of categorical values
            "some_other_factor": [1, 2],
        })

    run(
        PBTBenchmarkExample,
        name="pbt_test",
        scheduler=pbt,
        reuse_actors=True,
        verbose=False,
        stop={
            "training_iteration": 2000,
        },
        num_samples=4,
        config={
            "lr": 0.0001,
            # note: this parameter is perturbed but has no effect on
            # the model training in this example
            "some_other_factor": 1,
        })

variant_spec = example_module.get_variant_spec(example_args)
    trainable_class = example_module.get_trainable_class(example_args)

    experiment_kwargs = generate_experiment_kwargs(variant_spec, example_args)

    redis_address = ray.services.get_node_ip_address() + ':6379'

    ray.init(
        redis_address=redis_address,
        num_cpus=example_args.cpus,
        num_gpus=example_args.gpus,
        local_mode=False,
        include_webui=example_args.include_webui,
        temp_dir=example_args.temp_dir)

    tune.run(
        trainable_class,
        **experiment_kwargs,
        with_server=example_args.with_server,
        server_port=example_args.server_port,
        scheduler=None,
        queue_trials=True,
        reuse_actors=True)

ray.init()
    if args.scheduler == "fifo":
        sched = FIFOScheduler()
    elif args.scheduler == "asynchyperband":
        sched = AsyncHyperBandScheduler(
            time_attr="training_iteration",
            metric="mean_loss",
            mode="min",
            max_t=400,
            grace_period=60)
    else:
        raise NotImplementedError
    tune.register_trainable(
        "TRAIN_FN",
        lambda config, reporter: train_cifar10(args, config, reporter))
    tune.run(
        "TRAIN_FN",
        name=args.expname,
        verbose=2,
        scheduler=sched,
        stop={
            "mean_accuracy": 0.98,
            "training_iteration": 1 if args.smoke_test else args.epochs
        },
        resources_per_trial={
            "cpu": int(args.num_workers),
            "gpu": int(args.num_gpus)
        },
        num_samples=1 if args.smoke_test else args.num_samples,
        config={
            "lr": tune.sample_from(
                lambda spec: np.power(10.0, np.random.uniform(-4, -1))),

def run(model, model_args, result_dir, nmaxepochs):
    experiment = cifar10_experiment()
    try:
        with open('../config/redis_address', 'r') as f:
            address = f.read().strip()
            ray.init(redis_address=address)
    except:
        ray.init()
    ahb = AsyncHyperBandScheduler(reward_attr='mean_accuracy', max_t=nmaxepochs)
    trials = ray.tune.run(experiment, scheduler=ahb, raise_on_failed_trial=False, queue_trials=True)
    # trials = ray.tune.run(experiment, raise_on_failed_trial=False, queue_trials=True)
    trials = [trial for trial in trials if trial.last_result is not None]
    accuracy = [trial.last_result.get('mean_accuracy', float('-inf')) for trial in trials]

    checkpoint_path = Path(result_dir) / experiment.name
    checkpoint_path.mkdir(parents=True, exist_ok=True)
    checkpoint_path /= 'trial.pkl'
    with checkpoint_path.open('wb') as f:
        pickle.dump(trials, f)

    ex.add_artifact(str(checkpoint_path))
    return model, model_args, max(accuracy)

"optimizer_creator": optimizer_creator,
        "loss_creator": lambda config: nn.CrossEntropyLoss(),
        "train_function": train,
        "validation_function": validate,
        "num_replicas": num_replicas,
        "initialization_hook": initialization_hook,
        "use_gpu": use_gpu,
        "batch_size": 16 if test_mode else 512,
        "config": {
            "lr": tune.choice([1e-4, 1e-3, 5e-3, 1e-2]),
            "test_mode": test_mode
        },
        "backend": "nccl" if use_gpu else "gloo"
    }

    analysis = tune.run(
        PyTorchTrainable,
        num_samples=2,
        config=config,
        stop={"training_iteration": 2},
        verbose=2)

    return analysis.get_best_config(metric="mean_accuracy", mode="max")

config = {
            "model": {
                "custom_model": ["model1", "model2"][i % 2],
            },
            "gamma": random.choice([0.95, 0.99]),
        }
        return (None, obs_space, act_space, config)

    # Setup PPO with an ensemble of `num_policies` different policies
    policies = {
        "policy_{}".format(i): gen_policy(i)
        for i in range(args.num_policies)
    }
    policy_ids = list(policies.keys())

    tune.run(
        "PPO",
        stop={"training_iteration": args.num_iters},
        config={
            "env": "multi_cartpole",
            "log_level": "DEBUG",
            "simple_optimizer": args.simple,
            "num_sgd_iter": 10,
            "multiagent": {
                "policies": policies,
                "policy_mapping_fn": (
                    lambda agent_id: random.choice(policy_ids)),
            },

name="CCPPO",
    postprocess_fn=centralized_critic_postprocessing,
    loss_fn=loss_with_central_critic,
    before_loss_init=setup_mixins,
    grad_stats_fn=central_vf_stats,
    mixins=[
        LearningRateSchedule, EntropyCoeffSchedule, KLCoeffMixin,
        CentralizedValueMixin
    ])

CCTrainer = PPOTrainer.with_updates(name="CCPPOTrainer", default_policy=CCPPO)

if __name__ == "__main__":
    args = parser.parse_args()
    ModelCatalog.register_custom_model("cc_model", CentralizedCriticModel)
    tune.run(
        CCTrainer,
        stop={
            "timesteps_total": args.stop,
            "episode_reward_mean": 7.99,
        },
        config={
            "env": TwoStepGame,
            "batch_mode": "complete_episodes",
            "eager": False,
            "num_workers": 0,
            "multiagent": {
                "policies": {
                    "pol1": (None, Discrete(6), TwoStepGame.action_space, {}),
                    "pol2": (None, Discrete(6), TwoStepGame.action_space, {}),
                },
                "policy_mapping_fn": lambda x: "pol1" if x == 0 else "pol2",