diff --git a/torchx/runner/api.py b/torchx/runner/api.py index cd66bd489..e63f7a2fd 100644 --- a/torchx/runner/api.py +++ b/torchx/runner/api.py @@ -412,6 +412,14 @@ def dryrun( ): sched = self._scheduler(scheduler) resolved_cfg = sched.run_opts().resolve(cfg) + + # early validation before build workspace + with log_event( + "pre_build_validate", + scheduler, + ): + sched._pre_build_validate(app, scheduler, resolved_cfg) + if workspace and isinstance(sched, WorkspaceMixin): role = app.roles[0] old_img = role.image @@ -420,7 +428,13 @@ def dryrun( logger.info( 'To disable workspaces pass: --workspace="" from CLI or workspace=None programmatically.' ) - sched.build_workspace_and_update_role(role, workspace, resolved_cfg) + with log_event( + "build_workspace_and_update_role", + scheduler, + ) as ctx: + sched.build_workspace_and_update_role(role, workspace, resolved_cfg) + ctx._torchx_event.app_image = role.image + ctx._torchx_event.workspace = workspace if old_img != role.image: logger.info( diff --git a/torchx/schedulers/api.py b/torchx/schedulers/api.py index 3ef3c5899..6a9e1bf70 100644 --- a/torchx/schedulers/api.py +++ b/torchx/schedulers/api.py @@ -337,12 +337,19 @@ def log_iter( f"{self.__class__.__qualname__} does not support application log iteration" ) + def _pre_build_validate(self, app: AppDef, scheduler: str, cfg: T) -> None: + """ + validates before workspace build whether application is consistent with the scheduler. + + Raises error if application is not compatible with scheduler + """ + pass + def _validate(self, app: AppDef, scheduler: str, cfg: T) -> None: """ - Validates whether application is consistent with the scheduler. + Validates after workspace build whether application is consistent with the scheduler. - Raises: - ValueError: if application is not compatible with scheduler + Raises error if application is not compatible with scheduler """ for role in app.roles: if role.resource == NULL_RESOURCE: