diff --git a/nemo_run/run/experiment.py b/nemo_run/run/experiment.py index b213c490..643486e6 100644 --- a/nemo_run/run/experiment.py +++ b/nemo_run/run/experiment.py @@ -302,6 +302,7 @@ def __init__( jobs: list[Job | JobGroup] | None = None, base_dir: str | None = None, clean_mode: bool = False, + enable_goodbye_message: bool = True, ) -> None: """ Initializes an experiment run by creating its metadata directory and saving the experiment config. @@ -317,6 +318,7 @@ def __init__( _reconstruct: Generally, the user does not need to specify this flag. This is only set to True when using run.Experiment.from_dir. clean_mode: If True, disables all console output (logs, progress bars, etc.). Defaults to False. + enable_goodbye_message: if True, prints goodbye message after submitting job. Defaults to True. """ configure_logging(level=log_level) self._reconstruct = _reconstruct @@ -325,6 +327,7 @@ def __init__( self._title = title self._id = id or f"{title}_{int(time.time())}" + self._enable_goodbye_message = enable_goodbye_message base_dir = str(base_dir or get_nemorun_home()) self._exp_dir = os.path.join(base_dir, "experiments", title, self._id) @@ -1181,16 +1184,17 @@ def __exit__(self, exc_type, exc_value, tb): theme=os.environ.get("NEMO_RUN_CODE_THEME", "monokai"), ) ) - self.console.print( - Syntax( - self.GOODBYE_MESSAGE_BASH.format( - exp_id=self._id, - tasks=list(map(lambda job: job.id, self.jobs)), - ), - "shell", - theme=os.environ.get("NEMO_RUN_CODE_THEME", "monokai"), + if self._enable_goodbye_message: + self.console.print( + Syntax( + self.GOODBYE_MESSAGE_BASH.format( + exp_id=self._id, + tasks=list(map(lambda job: job.id, self.jobs)), + ), + "shell", + theme=os.environ.get("NEMO_RUN_CODE_THEME", "monokai"), + ) ) - ) def _repr_svg_(self): return self.to_config()._repr_svg_() diff --git a/nemo_run/run/torchx_backend/packaging.py b/nemo_run/run/torchx_backend/packaging.py index 093d190f..8df51dd5 100644 --- a/nemo_run/run/torchx_backend/packaging.py +++ b/nemo_run/run/torchx_backend/packaging.py @@ -138,7 +138,8 @@ def _get_details_from_script(fn_or_script: Script, serialize_configs: bool): fn_or_script, serialize_configs=True ) metadata = fn_or_script.metadata - env = env | fn_or_script.env + if fn_or_script.env: + env = env | fn_or_script.env launcher = executor.get_launcher() if executor.supports_launcher_transform():