diff --git a/dvc/api.py b/dvc/api.py index f9b2f4fe5..b7de3682e 100644 --- a/dvc/api.py +++ b/dvc/api.py @@ -121,12 +121,13 @@ def make_checkpoint(): import builtins from time import sleep + from dvc.env import DVC_CHECKPOINT, DVC_ROOT from dvc.stage.run import CHECKPOINT_SIGNAL_FILE - if os.getenv("DVC_CHECKPOINT") is None: + if os.getenv(DVC_CHECKPOINT) is None: return - root_dir = Repo.find_root() + root_dir = os.getenv(DVC_ROOT, Repo.find_root()) signal_file = os.path.join( root_dir, Repo.DVC_DIR, "tmp", CHECKPOINT_SIGNAL_FILE ) diff --git a/dvc/env.py b/dvc/env.py index 9b5ea67a8..d1b40671c 100644 --- a/dvc/env.py +++ b/dvc/env.py @@ -1,2 +1,4 @@ +DVC_CHECKPOINT = "DVC_CHECKPOINT" DVC_DAEMON = "DVC_DAEMON" DVC_PAGER = "DVC_PAGER" +DVC_ROOT = "DVC_ROOT" diff --git a/dvc/stage/run.py b/dvc/stage/run.py index 591350b64..68abc80c8 100644 --- a/dvc/stage/run.py +++ b/dvc/stage/run.py @@ -5,6 +5,7 @@ import subprocess import threading from contextlib import contextmanager +from dvc.env import DVC_CHECKPOINT, DVC_ROOT from dvc.utils import fix_env from .decorators import relock_repo, unlocked_repo @@ -47,7 +48,7 @@ def cmd_run(stage, *args, checkpoint_func=None, **kwargs): kwargs = {"cwd": stage.wdir, "env": fix_env(None), "close_fds": True} if checkpoint_func: # indicate that checkpoint cmd is being run inside DVC - kwargs["env"].update({"DVC_CHECKPOINT": "1"}) + kwargs["env"].update(_checkpoint_env(stage)) if os.name == "nt": kwargs["shell"] = True @@ -118,6 +119,10 @@ def run_stage(stage, dry=False, force=False, checkpoint_func=None, **kwargs): cmd_run(stage, checkpoint_func=checkpoint_func) +def _checkpoint_env(stage): + return {DVC_CHECKPOINT: "1", DVC_ROOT: stage.repo.root_dir} + + @contextmanager def checkpoint_monitor(stage, callback_func, proc, killed): if not callback_func: diff --git a/tests/func/experiments/test_checkpoints.py b/tests/func/experiments/test_checkpoints.py index 3d1c2d33c..503fcdac1 100644 --- a/tests/func/experiments/test_checkpoints.py +++ b/tests/func/experiments/test_checkpoints.py @@ -4,6 +4,7 @@ from textwrap import dedent import pytest from funcy import first +import dvc as dvc_module from dvc.exceptions import DvcException from dvc.repo.experiments import Experiments @@ -62,10 +63,15 @@ def checkpoint_stage(tmp_dir, scm, dvc): def test_new_checkpoint(tmp_dir, scm, dvc, checkpoint_stage, mocker): + from dvc.env import DVC_CHECKPOINT, DVC_ROOT + new_mock = mocker.spy(dvc.experiments, "new") + env_mock = mocker.spy(dvc_module.stage.run, "_checkpoint_env") dvc.experiments.run(checkpoint_stage.addressing, params=["foo=2"]) new_mock.assert_called_once() + env_mock.assert_called_once() + assert set(env_mock.return_value.keys()) == {DVC_CHECKPOINT, DVC_ROOT} assert (tmp_dir / "foo").read_text() == "5" assert ( tmp_dir / ".dvc" / "experiments" / "metrics.yaml"