greta-dev / greta-dev/greta

Give compute_options a structured object, so it can carry more than a device

Open
#839 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
607
Forks
67
Avg merge
3d 8h
Merged PRs (30d)
1

Description

## What I expect

That `compute_options` can express TensorFlow settings that are not "which
device", since several of them affect results rather than just placement.

## What happens instead

[`gpu_only()`](https://github.com/greta-dev/greta/blob/280dc906/R/utils.R#L863-L865) and
[`cpu_only()`](https://github.com/greta-dev/greta/blob/280dc906/R/utils.R#L869-L871) return the bare strings `"GPU"` and
`"CPU"`, which go straight to `tf$device()`. There is nowhere to put anything
else.

Two things want a home there.

**1. Op determinism on GPU.** Seeding fixes the random numbers, not the
arithmetic: some TensorFlow ops accumulate in a non-deterministic order on GPU,
so a seeded run can still vary. TensorFlow offers
`tf.config.experimental.enable_op_determinism()`, at a performance cost, which
is exactly the sort of trade-off a user should opt into rather than have chosen
for them. After #285, `set.seed()` makes `mcmc()` reproducible on CPU; on GPU it
gets you most of the way and no further, with nothing to reach for.

**2. `calculate()` disables the GPU for the whole session.** At
[`calculate.R#L217-L220`](https://github.com/greta-dev/greta/blob/280dc906/R/calculate.R#L217-L220) and
[`#L224-L227`](https://github.com/greta-dev/greta/blob/280dc906/R/calculate.R#L224-L227) it calls
`tensorflow::set_random_seed(disable_gpu = is_using_cpu(compute_options))`. That
function sets `CUDA_VISIBLE_DEVICES = -1` and **never puts it back**, so a single
`calculate(nsim = )` on the default `cpu_only()` hides the GPU from everything
later in the session:

```r
Sys.getenv("CUDA_VISIBLE_DEVICES") # ""
calculate(x, nsim = 1) # default cpu_only()
Sys.getenv("CUDA_VISIBLE_DEVICES") # "-1" - GPU now invisible
mcmc(m, compute_options = gpu_only()) # silently runs on CPU
```

`mcmc()` had the same problem briefly while fixing #285 and now seeds the Python
side directly instead. `calculate()` still needs a proper answer, and "which
device did the user actually ask for, and may we turn the GPU off" is
information that belongs in `compute_options`.

## Fix

Make `cpu_only()` and `gpu_only()` return a structured object rather than a
string, with arguments for the settings above — something like
`gpu_only(enable_op_determinism = TRUE)`.

Consumers that currently treat it as a string need updating:
`tf$device(compute_options)`, `compute_options == "CPU"`, `is_using_cpu()`,
`compute_text()`.

## This is not a breaking change

The constructor functions exist precisely so this could be added later, and the
string form was never part of the documented interface — `compute_options = "CPU"`
appears nowhere in `R/`, `tests/`, `vignettes/` or `inst/`. The `@param` has
promised this since it was written:

> *"Default is to use CPU only with `cpu_only()`. Use `gpu_only()` to use only
> GPU. [**In the future we will add more options for specifying CPU and GPU
> use.**](https://github.com/greta-dev/greta/blob/280dc906/R/inference.R#L42-L44)"*

## Related

Split out of #285, which is scoped to making `mcmc()` respect `set.seed()`.

Contributor guide

Open the contributing guide

Research direction

Start with gpu_only() and cpu_only() in R/utils.R, then trace compute_options through R/calculate.R and the listed consumers: tf$device(), is_using_cpu(), compute_text(), and the equality check against "CPU". Define the structured options around device selection and op determinism, update each consumer, and verify that calculate() does not leave the GPU disabled for later mcmc() calls.

Written by the indexing model from the issue text.

Assessment

Tech stack
r, tensorflow
Domain
machine-learning
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.