lattice-addons-pytorch 0.1.1
Lattice Add-ons
Requires Python: >=3.7

# State Needed from Agent
The addons use a set of environment variables that are passed from the agent to understand where they should write checkpoints and how often.
It utilizes the following:

`LATTICE_CHECKPOINT_ENDPOINT`

This is a string which indicates where the add-ons should write and look for checkpoints.
If using S3 for example, it takes the following format: `s3://breeze-checkpoint/path/to/storage/`
It will then use this endpoint to find the S3 bucket and the storage path and save the checkpoints there.

`LATTICE_RUN_ID`

This is a string which indicates the current rendezvous ID.
This is taken from the agent's `LATTICE_RDZV_ID` so that when storing checkpoints they can be stored
safely without conflicting with checkpoints stored by other jobs using the same storage backend.
Continuing the example above, the worker would store its checkpoints at
`s3://breeze-checkpoint/path/to/storage/<run_id>/checkpoint/`.
TODO: Make the usage of `run_id` and `rdzv_id` consistent.

`LATTICE_STEPS_PER_CHECKPOINT`

This is passed to the agent to tell it how often to write checkpoints.
This can be configured by the agent's --steps-per-checkpoint CLI arg or environment variable.


## How to pass state
There are two ways to pass state to the worker:

1. Add it to the `extra_env` dictionary seen [here](https://github.com/breezeml/lattice-agent/blob/8571870de3e2635819edeb4f48b289c9e0eb3b01/src/lattice/torch/launcher/api.py#L227)

2. Add it to the `worker_env` dictionary seen [here](https://github.com/breezeml/lattice-agent/blob/8571870de3e2635819edeb4f48b289c9e0eb3b01/src/lattice/elastic/agent/server/lattice_agent.py#L58)

    **Warning** Make sure all state passed through this dict is a string!
# Lattice PyTorch Addons

Lattice PyTorch Addons provides a convenient way of making your code resilient and elastic through
a simple function decorator.

## How to use

Include this import at the top of your training script.
```python
import lattice_addons.pytorch as lte
```

Initialize all objects needed for training before the main training loop, such as your model,
optimizer, etc.

```python
model = ...
optimizer = ...
loss_fn = ...
trainset, testset = get_datasets()

train_dataloader = ...
test_dataloader = ...
```

Once you have this done, make sure to encapsulate your full training loop in a single function, such as
the following:
```python
@lte.run
def full_train(model: torch.nn.Module, sampler, loader, optimizer: torch.optim.Optimizer, lr_scheduler, num_epochs: int, loss_fn, test_loader):
    for e in range(0, num_epochs):
        print("Starting epoch", e)

        run_train(e, model, optimizer, loss_fn, loader)

        run_eval(model, loss_fn, test_loader)
```
Note a few things about this function:

1. It takes all the objects required for training

2. It then runs the required number of epochs and training steps

3. **Decorator:** It is decorated by the `@lte.run` decorator

### `@lte.run` Decorator
This decorator creates wrappers of each component used for training in PyTorch.
It transparently forwards calls to the underlying object so that custom calls implemented
by the user will not be affected.
It then uses the standard calls to `forward()` in PyTorch to monitor training progress
and periodically save checkpoints.