autogluon / autogluon/autogluon

[timeseries] Add `num_cpus`, `num_gpus` parameter to `TimeSeriesPredictor.fit()` and `predict()` methods

Open
#4,220 1 comment 2 reactions 1 assignee Claimed by @shchur View on GitHub
enhancement module: timeseries
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.