From e58fcd864bbc7939f3bed1570df257af9b4263f3 Mon Sep 17 00:00:00 2001 From: Saugat Pachhai Date: Sat, 2 Oct 2021 16:12:44 +0545 Subject: [PATCH] [WIP] exp init: introduce `--with={default,live}` flag (#6716) * exp init: introduce `--with={default,live}` flag This flag will select the different mode of outputs of stage. With `--with=live`, it'll enable `live` setting in the stage. On interactive mode, it'll always ask user if they want to use live (unless `--with=live` is specified), and then it'll skip metrics/plots question and moves on to ask for path to dvclive path. Also, if live setting is enabled in the stage, it'll set the `models` to be a checkpoint output. * s/with/type and don't ask for live interactively --- dvc/command/experiments.py | 94 ++++++++++++++++++++++++++------------ 1 file changed, 64 insertions(+), 30 deletions(-) diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py index 98a79306f..0b0a6fdb4 100644 --- a/dvc/command/experiments.py +++ b/dvc/command/experiments.py @@ -837,28 +837,9 @@ class CmdExperimentsInit(CmdBase): DEFAULT_PARAMS = "params.yaml" PLOTS = "plots" DVCLIVE = "dvclive" - DEFAULT_NAME = "default" - @post_processing(dict) - def init_interactive(self, defaults=None, show_heading: bool = False): + def _prompt(self, prompts, defaults=None): defaults = defaults or {} - prompts = { - "cmd": "[b]Command[/b] to execute", - "code": "Path to a [b]code[/b] file/directory", - "data": "Path to a [b]data[/b] file/directory", - "models": "Path to a [b]model[/b] file/directory", - "metrics": "Path to a [b]metrics[/b] file", - "params": "Path to a [b]parameters[/b] file", - "plots": "Path to a [b]plots[/b] file/directory", - "live": "Path to log [b]dvclive[/b] outputs", - } - message = ( - "This command will guide you to set up your first stage in " - "[green]dvc.yaml[/green].\n" - ) - if show_heading: - ui.error_write(message, styled=True) - for key, prompt in prompts.items(): if key == "cmd": prompt_cls = RequiredPrompt @@ -868,14 +849,52 @@ class CmdExperimentsInit(CmdBase): value = prompt_cls.ask(prompt, console=ui.error_console, **kwargs) yield key, value + @post_processing(dict) + def init_interactive( + self, + defaults=None, + show_heading: bool = False, + live: bool = False, + ): + message = ( + "This command will guide you to set up your first stage in " + "[green]dvc.yaml[/green].\n" + ) + if show_heading: + ui.error_write(message, styled=True) + + yield from self._prompt( + { + "cmd": "[b]Command[/b] to execute", + "code": "Path to a [b]code[/b] file/directory", + "data": "Path to a [b]data[/b] file/directory", + "models": "Path to a [b]model[/b] file/directory", + "params": "Path to a [b]parameters[/b] file", + }, + defaults=defaults, + ) + + if not live: + yield from self._prompt( + { + "metrics": "Path to a [b]metrics[/b] file", + "plots": "Path to a [b]plots[/b] file/directory", + }, + defaults=defaults, + ) + return + + yield from self._prompt( + {"live": "Path to log [b]dvclive[/b] outputs"}, + defaults=defaults, + ) + def run(self): from dvc.command.stage import parse_cmd cmd = parse_cmd(self.args.cmd) if not self.args.interactive and not cmd: raise InvalidArgumentError("command is not specified") - if self.args.template: - raise NotImplementedError("template is not supported yet.") from dvc.dvcfile import make_dvcfile @@ -886,10 +905,11 @@ class CmdExperimentsInit(CmdBase): "metrics": self.DEFAULT_METRICS, "params": self.DEFAULT_PARAMS, "plots": self.PLOTS, + "live": self.DVCLIVE, } dvcfile = make_dvcfile(self.repo, "dvc.yaml") - name = self.args.name or self.DEFAULT_NAME + name = self.args.name or self.args.type dvcfile_exists = dvcfile.exists() if not self.args.force and dvcfile_exists and name in dvcfile.stages: @@ -905,16 +925,26 @@ class CmdExperimentsInit(CmdBase): config = {} # TODO context.maps.extend([config, global_defaults]) + with_live = self.args.type == "live" if self.args.interactive: - defaults = context.new_child({"live": self.DVCLIVE}) try: context = self.init_interactive( - defaults=defaults, show_heading=not dvcfile_exists + defaults=context, + show_heading=not dvcfile_exists, + live=with_live, ) except (KeyboardInterrupt, EOFError): ui.error_write() raise + elif with_live: + # suppress `metrics`/`params` if live is selected, unless + # also provided via cli, also make output to be a checkpoint. + context = context.new_child({"metrics": None, "params": None}) else: + # suppress live otherwise + context = context.new_child({"live": None}) + + if not self.args.interactive: d = compact( { "cmd": cmd, @@ -939,23 +969,24 @@ class CmdExperimentsInit(CmdBase): live = context.get("live") params_kv = [] - if "params" in context: + if context.get("params"): from dvc.utils.serialize import LOADERS path = context["params"] _, ext = os.path.splitext(path) params_kv = [{path: list(LOADERS[ext](path))}] + checkpoint_out = bool(context.get("live")) stage = self.repo.stage.add( name=name, cmd=command, deps=compact([code, data]), - outs=compact([models]), params=params_kv, metrics_no_cache=compact([metrics]), plots_no_cache=compact([plots]), live=live, force=self.args.force, + **{"checkpoints" if checkpoint_out else "outs": compact([models])}, ) if self.args.run: @@ -1506,9 +1537,6 @@ def add_parser(subparsers, parent_parser): action="store_true", help="Prompt for values that are not provided", ) - experiments_init_parser.add_argument( - "--template", help="Stage template to use to fill with provided values" - ) experiments_init_parser.add_argument( "-f", "--force", @@ -1553,6 +1581,12 @@ def add_parser(subparsers, parent_parser): experiments_init_parser.add_argument( "--live", help="Path to log dvclive outputs for your experiments" ) + experiments_init_parser.add_argument( + "--type", + choices=["default", "live"], + default="default", + help="Select type of stage to create", + ) experiments_init_parser.set_defaults(func=CmdExperimentsInit)