Construct staging and output manifests from FileStager and ResultsCollector

At the moment, JSON staging and outputs manifests are constructed by tracking all actions mapped by the `FileActionMapper` using a list `FileActionMapper.actions`. This makes the `FileActionMapper` stateful, requires including `file_type` as keyword argument for `BaseAction` and its children, requires defining a finalize()` method for `FileActionMapper` and for `BaseAction` and its children.

Paying the small price of refactoring `JsonTransferAction`, generate the staging manifest from `FileStager.transfer_tracker.remote_staging_actions` and the output manifest as `ResultsCollector` collects the outputs.

Author: kysrpex

pulsar/client/action_mapper.py

@@ -21,7 +21,6 @@
     Any,
     Dict,
     List,
-    Optional,
     Type,
 )
 from urllib.parse import urlencode
@@ -191,9 +190,6 @@ def __init__(self, client=None, config=None):
         self.ssh_port = config.get("ssh_port", None)
         self.mappers = mappers_from_dicts(config.get("paths", []))
         self.files_endpoint = config.get("files_endpoint", None)
-        self.actions = []
-        # Might want to make the working directory available here so that we know where to place archive
-        # for archive action
 
     def action(self, source, type, mapper=None):
         path = source.get("path", None)
@@ -204,14 +200,10 @@ def action(self, source, type, mapper=None):
         if mapper:
             file_lister = mapper.file_lister
             action_kwds = mapper.action_kwds
-        action = action_class(source, file_lister=file_lister, file_type=type, **action_kwds)
+        action = action_class(source, file_lister=file_lister, **action_kwds)
         self.__process_action(action, type)
-        self.actions.append(action)
         return action
 
-    def finalize(self):
-        return [_ for _ in (action.finalize() for action in self.actions) if _]
-
     def unstructured_mappers(self):
         """ Return mappers that will map 'unstructured' files (i.e. go beyond
         mapping inputs, outputs, and config files).
@@ -273,7 +265,6 @@ def __process_action(self, action, file_type):
         """ Extension point to populate extra action information after an
         action has been created.
         """
-        action.file_type = file_type
         if getattr(action, "inject_url", False):
             self.__inject_url(action, file_type)
         if getattr(action, "inject_ssh_properties", False):
@@ -309,12 +300,10 @@ class BaseAction:
     whole_directory_transfer_supported = False
     action_spec: Dict[str, Any] = {}
     action_type: str
-    file_type: Optional[str] = None
 
-    def __init__(self, source, file_lister=None, file_type=None):
+    def __init__(self, source, file_lister=None):
         self.source = source
         self.file_lister = file_lister or DEFAULT_FILE_LISTER
-        self.file_type = file_type
 
     @property
     def path(self):
@@ -353,9 +342,6 @@ def _extend_base_dict(self, **kwds):
         base_dict.update(**kwds)
         return base_dict
 
-    def finalize(self):
-        pass
-
     def to_dict(self):
         return self._extend_base_dict()
 
@@ -404,8 +390,8 @@ class RewriteAction(BaseAction):
     action_type = "rewrite"
     staging = STAGING_ACTION_NONE
 
-    def __init__(self, source, file_lister=None, source_directory=None, destination_directory=None, file_type=None):
-        super().__init__(source, file_lister=file_lister, file_type=file_type)
+    def __init__(self, source, file_lister=None, source_directory=None, destination_directory=None):
+        super().__init__(source, file_lister=file_lister)
         self.source_directory = source_directory
         self.destination_directory = destination_directory
 
@@ -481,8 +467,8 @@ class RemoteTransferAction(BaseAction):
     action_type = "remote_transfer"
     staging = STAGING_ACTION_REMOTE
 
-    def __init__(self, source, file_lister=None, url=None, file_type=None):
-        super().__init__(source, file_lister=file_lister, file_type=file_type)
+    def __init__(self, source, file_lister=None, url=None):
+        super().__init__(source, file_lister=file_lister)
         self.url = url
 
     def to_dict(self):
@@ -509,8 +495,8 @@ class RemoteTransferTusAction(BaseAction):
     action_type = "remote_transfer_tus"
     staging = STAGING_ACTION_REMOTE
 
-    def __init__(self, source, file_lister=None, url=None, file_type=None):
-        super().__init__(source, file_lister=file_lister, file_type=file_type)
+    def __init__(self, source, file_lister=None, url=None):
+        super().__init__(source, file_lister=file_lister)
         self.url = url
 
     def to_dict(self):
@@ -537,30 +523,39 @@ class JsonTransferAction(BaseAction):
     action_type = "json_transfer"
     staging = STAGING_ACTION_REMOTE
 
-    def __init__(self, source, file_lister=None, url=None, file_type=None):
-        super().__init__(source, file_lister, file_type)
+    def __init__(self, source, file_lister=None, url=None, from_path=None, to_path=None):
+        super().__init__(source, file_lister)
         self.url = url
-        self._from_path = None
-        self._to_path = None
+
+        self._from_path = from_path
+        self._to_path = to_path
+        # `from_path` and `to_path` are mutually exclusive, only one of them should be set
 
     @classmethod
     def from_dict(cls, action_dict):
-        return JsonTransferAction(source=action_dict["source"], url=action_dict["url"])
+        return JsonTransferAction(
+            source=action_dict["source"],
+            url=action_dict["url"],
+            from_path=action_dict.get("from_path"),
+            to_path=action_dict.get("to_path")
+        )
 
     def to_dict(self):
-        return self._extend_base_dict(url=self.url)
+        return self._extend_base_dict(**self.to_staging_manifest_entry())
 
     def write_to_path(self, path):
-        self._to_path = path
+        self._from_path, self._to_path = None, path
 
     def write_from_path(self, pulsar_path: str):
-        self._from_path = pulsar_path
+        self._from_path, self._to_path = pulsar_path, None
 
-    def finalize(self):
+    def to_staging_manifest_entry(self):
+        staging_manifest_entry = dict(url=self.url)
+        if self._from_path:
+            staging_manifest_entry["from_path"] = self._from_path
         if self._to_path:
-            return {"url": self.url, "to_path": self._to_path}
-        else:
-            return {"url": self.url, "from_path": self._from_path}
+            staging_manifest_entry["to_path"] = self._to_path
+        return staging_manifest_entry
 
 
 class RemoteObjectStoreCopyAction(BaseAction):
@@ -911,6 +906,7 @@ def unstructured_map(self, path):
 
 __all__ = (
     'FileActionMapper',
+    'JsonTransferAction',
     'path_type',
     'from_dict',
     'MessageAction',

pulsar/client/staging/down.py

@@ -3,12 +3,13 @@
 from contextlib import contextmanager
 from json import loads
 from logging import getLogger
+from typing import Optional
 from os.path import (
     join,
     relpath,
 )
 
-from ..action_mapper import FileActionMapper
+from ..action_mapper import FileActionMapper, JsonTransferAction
 from ..staging import COMMAND_VERSION_FILENAME
 
 log = getLogger(__name__)
@@ -64,16 +65,17 @@ def __init__(self, output_collector, action_mapper, client_outputs, pulsar_outpu
         self.working_directory_contents = pulsar_outputs.working_directory_contents or []
         self.metadata_directory_contents = pulsar_outputs.metadata_directory_contents or []
         self.job_directory_contents = pulsar_outputs.job_directory_contents or []
+        self.output_manifest: Optional[list] = None
 
     def collect(self):
+        self.output_manifest = []
+
         self.__collect_working_directory_outputs()
         self.__collect_outputs()
         self.__collect_version_file()
         self.__collect_other_working_directory_files()
         self.__collect_metadata_directory_files()
         self.__collect_job_directory_files()
-        # Give actions that require a final action, like those that write a manifest, to write out their content
-        self.__finalize_action_mapper()
         # finalize collection here for executors that need this ?
         return self.exception_tracker.collection_failure_exceptions
 
@@ -137,9 +139,6 @@ def __collect_job_directory_files(self):
             'output_jobdir',
         )
 
-    def __finalize_action_mapper(self):
-        self.action_mapper.finalize()
-
     def __realized_dynamic_file_source_references(self):
         references = {"filename": [], "extra_files": []}
 
@@ -212,7 +211,10 @@ def _attempt_collect_output(self, output_type, path, name=None):
     def _collect_output(self, output_type, action, name):
         log.info("collecting output {} with action {}".format(name, action))
         try:
-            return self.output_collector.collect_output(self, output_type, action, name)
+            collect_result = self.output_collector.collect_output(self, output_type, action, name)
+            if isinstance(action, JsonTransferAction):
+                self.output_manifest.append(action.to_staging_manifest_entry())
+            return collect_result
         except Exception as e:
             if _allow_collect_failure(output_type):
                 log.warning(

pulsar/client/staging/up.py

@@ -16,6 +16,7 @@
 
 from ..action_mapper import (
     FileActionMapper,
+    JsonTransferAction,
     MessageAction,
     path_type,
 )
@@ -72,15 +73,19 @@ def submit_job(client, client_job_description, job_config=None):
     launch_kwds["dynamic_file_sources"] = client_job_description.client_outputs.dynamic_file_sources
     launch_kwds["token_endpoint"] = client.token_endpoint
 
-    # populate `to_path`
+    # generate staging manifest
     staging_manifest = []
-    for action in file_stager.action_mapper.actions:
-        if action.file_type not in ("output", "output_workdir"):
+    for action_description in remote_staging_actions:
+        action_dict = action_description["action"]
+        is_json_transfer_action = action_dict.get("action_type") == JsonTransferAction.action_type
+        is_not_output_action = action_description.get("type") not in ("output", "output_workdir")
+        if is_json_transfer_action and is_not_output_action:
+            file_type = action_description.get("type")
+            action = JsonTransferAction.from_dict(action_dict)
             name = basename(action.path)
-            path = file_stager.job_directory.calculate_path(name, action.file_type)
+            path = file_stager.job_directory.calculate_path(name, file_type)
             action.write_to_path(path)
-            staging_manifest.append(action.finalize())
-
+            staging_manifest.append(action.to_staging_manifest_entry())
     if staging_manifest:
         launch_kwds["staging_manifest"] = staging_manifest
 

pulsar/managers/staging/post.py

@@ -20,7 +20,9 @@ def postprocess(job_directory, action_executor, was_cancelled):
             staging_config = job_directory.load_metadata("launch_config").get("remote_staging", None)
         else:
             staging_config = None
-        file_action_mapper, collected = _collect_outputs(job_directory, staging_config, action_executor, was_cancelled)
+        file_action_mapper, _, collected = _collect_outputs(
+            job_directory, staging_config, action_executor, was_cancelled
+        )
         return collected
     finally:
         job_directory.write_file("postprocessed", "")
@@ -29,6 +31,7 @@ def postprocess(job_directory, action_executor, was_cancelled):
 
 def _collect_outputs(job_directory, staging_config, action_executor, was_cancelled):
     collected = True
+    output_manifest = None
     if "action_mapper" in staging_config:
         file_action_mapper = action_mapper.FileActionMapper(config=staging_config["action_mapper"])
         client_outputs = staging.ClientOutputs.from_dict(staging_config["client_outputs"])
@@ -39,7 +42,8 @@ def _collect_outputs(job_directory, staging_config, action_executor, was_cancell
         if collection_failure_exceptions:
             log.warn("Failures collecting results %s" % collection_failure_exceptions)
             collected = False
-    return file_action_mapper, collected
+        output_manifest = results_collector.output_manifest
+    return file_action_mapper, output_manifest, collected
 
 
 def realized_dynamic_file_sources(job_directory):

pulsar/scripts/collect_output_manifest.py

@@ -23,15 +23,14 @@ def collect_outputs(job_directory: str, staging_config_path: str, output_manifes
     with open(staging_config_path) as staging_fh:
         staging_config = json.load(staging_fh)
 
-    action_mapper, _ = _collect_outputs(
+    action_mapper, output_manifest, _ = _collect_outputs(
         job_directory_,
         staging_config=staging_config,
         action_executor=RetryActionExecutor(),
         was_cancelled=lambda: False
     )
-    new_manifest = action_mapper.finalize()
     with open(output_manifest_path, "w") as manifest_fh:
-        json.dump(new_manifest, manifest_fh)
+        json.dump(output_manifest, manifest_fh)
 
 
 def main():