torchx support early validation before workspace build

Differential Revision: D69463377

Pull Request resolved: #1007

Author: tonykao8080

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(

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: