Repository: iterative/dvc

Issue:
exp run: fails with multiple checkpoint outs in a stage
# Bug Report

## Description

Might not be technically a bug, but it seems impossible right now to have multiple checkpoints. For example, if I have my model weights as one output and a file with the current step/epoch as another, I want both of these to act like checkpoints outputs.

### Reproduce

```console
(dvc) █▓▒░dave@davids-air█▓▒░ Thu May 13 02:45:44P
~/repo/ dvc stage add -v -n mystage -c foo -c bar "echo foo > foo; echo bar > bar"
2021-05-13 14:45:50,063 DEBUG: Check for update is enabled.
2021-05-13 14:45:50,091 DEBUG: Trying to spawn '['daemon', '-q', 'updater']'
2021-05-13 14:45:50,095 DEBUG: Spawned '['daemon', '-q', 'updater']'
Creating 'dvc.yaml'
Adding stage 'mystage' in 'dvc.yaml'
2021-05-13 14:45:50,622 DEBUG: Adding 'foo' to '.gitignore'.
2021-05-13 14:45:50,624 DEBUG: Adding 'bar' to '.gitignore'.

To track the changes with git, run:

        git add .gitignore dvc.yaml
2021-05-13 14:45:50,626 DEBUG: Analytics is disabled.
(dvc) █▓▒░dave@davids-air█▓▒░ Thu May 13 02:45:50P
~/repo/ dvc exp run -v
2021-05-13 14:45:59,372 DEBUG: Check for update is enabled.
2021-05-13 14:45:59,387 DEBUG: Detaching HEAD at 'HEAD'
2021-05-13 14:45:59,412 DEBUG: Stashing workspace
2021-05-13 14:45:59,418 DEBUG: No changes to stash
2021-05-13 14:45:59,447 DEBUG: Lockfile for 'dvc.yaml' not found
2021-05-13 14:45:59,457 DEBUG: Stashing changes in 'refs/exps/stash'
2021-05-13 14:45:59,483 DEBUG: Stashed experiment '148b77d' with baseline '9bf289d' for future execution.
2021-05-13 14:45:59,485 DEBUG: Restore HEAD to 'main'
2021-05-13 14:45:59,486 DEBUG: Detaching HEAD at '9bf289d9b899999d9d0141da0647f59f327cb2a2'
2021-05-13 14:45:59,488 DEBUG: Popping from stash 'refs/exps/stash'
2021-05-13 14:45:59,488 DEBUG: Applying stash commit '148b77d99e4d46c9f2efa9829fec2c4d0c4bb1fd'
2021-05-13 14:45:59,491 DEBUG: Dropping 'refs/exps/stash@{0}'
2021-05-13 14:45:59,502 DEBUG: Running repro in '/Users/dave/repo'
2021-05-13 14:45:59,503 DEBUG: Removing '/Users/dave/repo/.dvc/tmp/repro.dat'
2021-05-13 14:45:59,509 DEBUG: Lockfile for 'dvc.yaml' not found
2021-05-13 14:45:59,518 DEBUG: Lockfile for 'dvc.yaml' not found
2021-05-13 14:45:59,521 DEBUG: Lockfile for 'dvc.yaml' not found
2021-05-13 14:45:59,735 DEBUG: stage: 'mystage' changed.
Running stage 'mystage':
2021-05-13 14:45:59,740 DEBUG: Restore HEAD to 'main'
2021-05-13 14:45:59,741 ERROR: failed to reproduce 'dvc.yaml': Duplicated env variable
------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/dave/Code/dvc/dvc/repo/reproduce.py", line 196, in _reproduce_stages
    ret = _reproduce_stage(stage, **kwargs)
  File "/Users/dave/Code/dvc/dvc/repo/reproduce.py", line 39, in _reproduce_stage
    stage = stage.reproduce(**kwargs)
  File "/Users/dave/miniconda3/envs/dvc/lib/python3.9/site-packages/funcy/decorators.py", line 39, in wrapper
    return deco(call, *dargs, **dkwargs)
  File "/Users/dave/Code/dvc/dvc/stage/decorators.py", line 36, in rwlocked
    return call()
  File "/Users/dave/miniconda3/envs/dvc/lib/python3.9/site-packages/funcy/decorators.py", line 60, in __call__
    return self._func(*self._args, **self._kwargs)
  File "/Users/dave/Code/dvc/dvc/stage/__init__.py", line 407, in reproduce
    self.run(**kwargs)
  File "/Users/dave/miniconda3/envs/dvc/lib/python3.9/site-packages/funcy/decorators.py", line 39, in wrapper
    return deco(call, *dargs, **dkwargs)
  File "/Users/dave/Code/dvc/dvc/stage/decorators.py", line 36, in rwlocked
    return call()
  File "/Users/dave/miniconda3/envs/dvc/lib/python3.9/site-packages/funcy/decorators.py", line 60, in __call__
    return self._func(*self._args, **self._kwargs)
  File "/Users/dave/Code/dvc/dvc/stage/__init__.py", line 514, in run
    run_stage(self, dry, force, **kwargs)
  File "/Users/dave/Code/dvc/dvc/stage/run.py", line 160, in run_stage
    run(stage, dry=dry, checkpoint_func=checkpoint_func, run_env=run_env)
  File "/Users/dave/Code/dvc/dvc/stage/decorators.py", line 45, in wrapper
    ret = f(stage, *args, **kwargs)
  File "/Users/dave/Code/dvc/dvc/stage/run.py", line 131, in cmd_run
    kwargs = prepare_kwargs(
  File "/Users/dave/Code/dvc/dvc/stage/run.py", line 50, in prepare_kwargs
    kwargs["env"].update(stage.env(checkpoint_func=checkpoint_func))
  File "/Users/dave/Code/dvc/dvc/stage/__init__.py", line 300, in env
    raise DvcException("Duplicated env variable")
dvc.exceptions.DvcException: Duplicated env variable

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/dave/Code/dvc/dvc/main.py", line 55, in main
    ret = cmd.do_run()
  File "/Users/dave/Code/dvc/dvc/command/base.py", line 50, in do_run
    return self.run()
  File "/Users/dave/Code/dvc/dvc/command/experiments.py", line 549, in run
    results = self.repo.experiments.run(
  File "/Users/dave/Code/dvc/dvc/repo/experiments/__init__.py", line 889, in run
    return run(self.repo, *args, **kwargs)
  File "/Users/dave/Code/dvc/dvc/repo/__init__.py", line 49, in wrapper
    return f(repo, *args, **kwargs)
  File "/Users/dave/Code/dvc/dvc/repo/experiments/run.py", line 32, in run
    return repo.experiments.reproduce_one(
  File "/Users/dave/Code/dvc/dvc/repo/experiments/__init__.py", line 421, in reproduce_one
    results = self._workspace_repro()
  File "/Users/dave/Code/dvc/dvc/repo/experiments/__init__.py", line 57, in wrapper
    ret = f(exp, *args, **kwargs)
  File "/Users/dave/Code/dvc/dvc/repo/experiments/__init__.py", line 771, in _workspace_repro
    exec_result = BaseExecutor.reproduce(
  File "/Users/dave/Code/dvc/dvc/repo/experiments/executor/base.py", line 294, in reproduce
    stages = dvc_reproduce(
  File "/Users/dave/Code/dvc/dvc/repo/__init__.py", line 49, in wrapper
    return f(repo, *args, **kwargs)
  File "/Users/dave/Code/dvc/dvc/repo/scm_context.py", line 14, in run
    return method(repo, *args, **kw)
  File "/Users/dave/Code/dvc/dvc/repo/reproduce.py", line 135, in reproduce
    return _reproduce_stages(self.graph, list(stages), **kwargs)
  File "/Users/dave/Code/dvc/dvc/repo/reproduce.py", line 213, in _reproduce_stages
    raise ReproductionError(stage.relpath) from exc
dvc.exceptions.ReproductionError: failed to reproduce 'dvc.yaml'
------------------------------------------------------------
2021-05-13 14:45:59,747 DEBUG: Analytics is disabled.
```

### Expected

I would expect stages to support any number of checkpoints outputs.

### Environment information

**Output of `dvc doctor`:**

```console
$ dvc doctor
DVC version: 2.1.0+81f42a.mod
---------------------------------
Platform: Python 3.9.1 on macOS-10.16-x86_64-i386-64bit
Supports: gdrive, gs, hdfs, webhdfs, http, https, ssh, oss, webdav, webdavs
Cache types: <https://error.dvc.org/no-dvc-cache>
Caches: local
Remotes: local
Workspace directory: apfs on /dev/disk3s1s1
Repo: dvc, git
```


Discussion:
This is a bug in how we set up our stage execution environment when we have checkpoints. The intended behavior for experiments is that you should be able to have multiple checkpoint outs in a given stage

Resolving patch:
diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py
--- a/dvc/stage/__init__.py
+++ b/dvc/stage/__init__.py
@@ -296,8 +296,11 @@ def env(self, checkpoint_func=None) -> Env:
 
         for out in self.outs:
             current = self._read_env(out, checkpoint_func=checkpoint_func)
-            if set(env.keys()).intersection(set(current.keys())):
-                raise DvcException("Duplicated env variable")
+            if any(
+                env.get(key) != current.get(key)
+                for key in set(env.keys()).intersection(current.keys())
+            ):
+                raise DvcException("Conflicting values for env variable")
             env.update(current)
         return env
 


Test patch:
diff --git a/tests/unit/stage/test_stage.py b/tests/unit/stage/test_stage.py
--- a/tests/unit/stage/test_stage.py
+++ b/tests/unit/stage/test_stage.py
@@ -122,3 +122,23 @@ def test_external_outs(tmp_path_factory, dvc):
 
     create_stage(Stage, dvc, "path.dvc", outs=["remote://myremote/foo"])
     create_stage(Stage, dvc, "path.dvc", outs=[os.fspath(foo)], external=True)
+
+
+def test_env(dvc):
+    from dvc.env import DVC_ROOT
+    from dvc.exceptions import DvcException
+    from dvc.stage import create_stage
+
+    stage = create_stage(Stage, dvc, "path.dvc", outs=["foo", "bar"])
+
+    with mock.patch.object(stage, "_read_env", return_value={"foo": "foo"}):
+        env = stage.env()
+        assert env == {DVC_ROOT: dvc.root_dir, "foo": "foo"}
+
+    def mock_read_env(out, **kwargs):
+        return {"foo": str(out)}
+
+    with pytest.raises(DvcException) as exc:
+        with mock.patch.object(stage, "_read_env", mock_read_env):
+            env = stage.env()
+        assert exc.value == "Conflicting values for env variable"