[Optim] Consider pickle for optimizer state file in run (with option to toggle)

[eddiebergman]

While refactoring meta-hyper, I wanted to benchmark why running the following without timeouts and a 0.0001 polling wasn't near instant on a single worker.

python neps_examples/basic_usage/hyperparameters.py

I used a tool call py-spy to generate a flamegraph (attached). As you can see, almost all of the time is spent in line 511, specifically the with sampler.using_state() with it's __enter__ and __exit__. This is all just yaml serialization of the optimizer state. For reference you can see the relative time of this serialization/desreialization relative to the acquisition sampling on the right side of the flamegraph.

I tried updating the contextmanager to use pickle just to check:

# Original, uses yaml for optimizer state
@contextmanager
def using_state(self, state_file: Path) -> Iterator[Sampler]:
    if state_file.exists():
        state = deserialize(state_file)
        self.load_state(state)

    yield self

    serialize(self.get_state(), path=state_file)

# With pickle, uses binary format
@contextmanager
def using_state(self, state_file: Path) -> Iterator[Sampler]:
    if state_file.exists():
        with state_file.open("rb") as f:
            state = pickle.load(f)
        self.load_state(state)

    yield self

    with state_file.open("wb") as f:
        pickle.dump(self.get_state(), f)

Here's the flamegraph and the using_state() doesn't even show up as taking anytime anymore. (It is there, it's one of the tiny tiny boxes on the left.

---

Not sure what we want to do with this information but if NePS gets benchmarked using a time constraint on hyperparameter benchmarks, its going to get smoked just due to serialization