autogluon / autogluon/autogluon
[timeseries] Add `num_cpus`, `num_gpus` parameter to `TimeSeriesPredictor.fit()` and `predict()` methods
- Dominant language
- Python
- Stars
- 10.7k
- Forks
- 1.2k
- Avg merge
- 21h 29m
- Merged PRs (30d)
- 57
Description
## Description
I would like to be able to limit the number of CPUs AutoGluon uses in training. Specifically I want to limit the cores used by the TFT model.
AutoGluon tabular has a parameter `TabularPredictor(...).fit(..., num_cpus = NUM_CORES_YOU_WANT)` https://github.com/autogluon/autogluon/blob/6583339731e92421824f14c5eb8a9997e1283c49/tabular/src/autogluon/tabular/predictor/predictor.py#L407
As far as I can tell, I have not been able to find the same in the TimeSeries predictor.
The reason for not wanting AG to decide for me is that I am running my training jobs in K8s, and it seems that AG's method for getting the max available CPUs is getting the total number of CPUs available in the node, but my specific pod is only allowed to use a fraction of that. This is leading to constant throttling in my training jobs.
The only model which has caused me issues is TFT, I have tried to limit the number of CPUs used by the `lightning.Trainer` but it does not seem to have changed anything. I will now try to use `torch.set_num_threads` but I still think it would be nice to be able to configure this manually, or somehow override the limits found by autogluon to match cloud workflows.
EDIT: the combination of
```yml
hyperparameters:
TemporalFusionTransformer:
trainer_kwargs:
devices: 2
```
and
```python
torch.set_num_threads(4)
```
Allows me to successfully control my CPU usage
Contributor guide
Assessment
This issue has not been assessed yet.