tracdap-runtime 0.7.0rc1__py3-none-any.whl → 0.8.0b2__py3-none-any.whl

tracdap/rt/_exec/functions.py

@@ -13,10 +13,10 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from __future__ import annotations
-
+import copy
 import datetime
 import abc
+import pathlib
 import random
 import dataclasses as dc  # noqa
 
@@ -28,6 +28,7 @@ import tracdap.rt._exec.graph_builder as _graph
 import tracdap.rt._impl.config_parser as _cfg_p  # noqa
 import tracdap.rt._impl.type_system as _types  # noqa
 import tracdap.rt._impl.data as _data  # noqa
+import tracdap.rt._impl.logging as _logging  # noqa
 import tracdap.rt._impl.storage as _storage  # noqa
 import tracdap.rt._impl.models as _models  # noqa
 import tracdap.rt._impl.util as _util  # noqa
@@ -227,11 +228,22 @@ class BuildJobResultFunc(NodeFunction[_config.JobResult]):
         job_result.jobId = self.node.job_id
         job_result.statusCode = meta.JobStatusCode.SUCCEEDED
 
+        if self.node.result_id is not None:
+
+            result_def = meta.ResultDefinition()
+            result_def.jobId = _util.selector_for(self.node.job_id)
+            result_def.statusCode = meta.JobStatusCode.SUCCEEDED
+
+            result_key = _util.object_key(self.node.result_id)
+            result_obj = meta.ObjectDefinition(objectType=meta.ObjectType.RESULT, result=result_def)
+
+            job_result.results[result_key] = result_obj
+
         # TODO: Handle individual failed results
 
-        for obj_id, node_id in self.node.outputs.objects.items():
+        for obj_key, node_id in self.node.outputs.objects.items():
             obj_def = _ctx_lookup(node_id, ctx)
-            job_result.results[obj_id] = obj_def
+            job_result.results[obj_key] = obj_def
 
         for bundle_id in self.node.outputs.bundles:
             bundle = _ctx_lookup(bundle_id, ctx)
@@ -241,9 +253,9 @@ class BuildJobResultFunc(NodeFunction[_config.JobResult]):
 
             runtime_outputs = _ctx_lookup(self.node.runtime_outputs, ctx)
 
-            for obj_id, node_id in runtime_outputs.objects.items():
+            for obj_key, node_id in runtime_outputs.objects.items():
                 obj_def = _ctx_lookup(node_id, ctx)
-                job_result.results[obj_id] = obj_def
+                job_result.results[obj_key] = obj_def
 
             for bundle_id in runtime_outputs.bundles:
                 bundle = _ctx_lookup(bundle_id, ctx)
@@ -252,37 +264,6 @@ class BuildJobResultFunc(NodeFunction[_config.JobResult]):
         return job_result
 
 
-class SaveJobResultFunc(NodeFunction[None]):
-
-    def __init__(self, node: SaveJobResultNode):
-        super().__init__()
-        self.node = node
-
-    def _execute(self, ctx: NodeContext) -> None:
-
-        job_result = _ctx_lookup(self.node.job_result_id, ctx)
-
-        if not self.node.result_spec.save_result:
-            return None
-
-        job_result_format = self.node.result_spec.result_format
-        job_result_str = _cfg_p.ConfigQuoter.quote(job_result, job_result_format)
-        job_result_bytes = bytes(job_result_str, "utf-8")
-
-        job_key = _util.object_key(job_result.jobId)
-        job_result_file = f"job_result_{job_key}.{self.node.result_spec.result_format}"
-        job_result_path = pathlib \
-            .Path(self.node.result_spec.result_dir) \
-            .joinpath(job_result_file)
-
-        _util.logger_for_object(self).info(f"Saving job result to [{job_result_path}]")
-
-        with open(job_result_path, "xb") as result_stream:
-            result_stream.write(job_result_bytes)
-
-        return None
-
-
 class DataViewFunc(NodeFunction[_data.DataView]):
 
     def __init__(self, node: DataViewNode):
@@ -296,8 +277,13 @@ class DataViewFunc(NodeFunction[_data.DataView]):
 
         # Map empty item -> emtpy view (for optional inputs not supplied)
         if root_item.is_empty():
-            return _data.DataView.create_empty()
+            return _data.DataView.create_empty(root_item.object_type)
 
+        # Handle file data views
+        if root_item.object_type == meta.ObjectType.FILE:
+            return _data.DataView.for_file_item(root_item)
+
+        # Everything else is a regular data view
         if self.node.schema is not None and len(self.node.schema.table.fields) > 0:
             trac_schema = self.node.schema
         else:
@@ -322,7 +308,11 @@ class DataItemFunc(NodeFunction[_data.DataItem]):
 
         # Map empty view -> emtpy item (for optional outputs not supplied)
         if data_view.is_empty():
-            return _data.DataItem.create_empty()
+            return _data.DataItem.create_empty(data_view.object_type)
+
+        # Handle file data views
+        if data_view.object_type == meta.ObjectType.FILE:
+            return data_view.file_item
 
         # TODO: Support selecting data item described by self.node
 
@@ -342,25 +332,24 @@ class DataResultFunc(NodeFunction[ObjectBundle]):
 
     def _execute(self, ctx: NodeContext) -> ObjectBundle:
 
-        data_item = _ctx_lookup(self.node.data_item_id, ctx)
+        data_spec = _ctx_lookup(self.node.data_save_id, ctx)
 
-        # Do not record output metadata for optional outputs that are empty
-        if data_item.is_empty():
-            return {}
+        result_bundle = dict()
 
-        data_spec = _ctx_lookup(self.node.data_spec_id, ctx)
+        # Do not record output metadata for optional outputs that are empty
+        if data_spec.is_empty():
+            return result_bundle
 
-        # TODO: Check result of save operation
-        # save_result = _ctx_lookup(self.node.data_save_id, ctx)
+        if self.node.data_key is not None:
+            result_bundle[self.node.data_key] = meta.ObjectDefinition(objectType=meta.ObjectType.DATA, data=data_spec.data_def)
 
-        data_result = meta.ObjectDefinition(objectType=meta.ObjectType.DATA, data=data_spec.data_def)
-        storage_result = meta.ObjectDefinition(objectType=meta.ObjectType.STORAGE, storage=data_spec.storage_def)
+        if self.node.file_key is not None:
+            result_bundle[self.node.file_key] = meta.ObjectDefinition(objectType=meta.ObjectType.FILE, file=data_spec.file_def)
 
-        bundle = {
-            self.node.data_key: data_result,
-            self.node.storage_key: storage_result}
+        if self.node.storage_key is not None:
+            result_bundle[self.node.storage_key] = meta.ObjectDefinition(objectType=meta.ObjectType.STORAGE, storage=data_spec.storage_def)
 
-        return bundle
+        return result_bundle
 
 
 class DynamicDataSpecFunc(NodeFunction[_data.DataSpec]):
@@ -443,11 +432,7 @@ class DynamicDataSpecFunc(NodeFunction[_data.DataSpec]):
 
         # Dynamic data def will always use an embedded schema (this is no ID for an external schema)
 
-        return _data.DataSpec(
-            data_item,
-            data_def,
-            storage_def,
-            schema_def=None)
+        return _data.DataSpec.create_data_spec(data_item, data_def, storage_def, schema_def=None)
 
 
 class _LoadSaveDataFunc(abc.ABC):
@@ -455,6 +440,16 @@ class _LoadSaveDataFunc(abc.ABC):
     def __init__(self, storage: _storage.StorageManager):
         self.storage = storage
 
+    @classmethod
+    def _choose_data_spec(cls, spec_id, spec, ctx: NodeContext):
+
+        if spec_id is not None:
+            return _ctx_lookup(spec_id, ctx)
+        elif spec is not None:
+            return spec
+        else:
+            raise _ex.EUnexpected()
+
     def _choose_copy(self, data_item: str, storage_def: meta.StorageDefinition) -> meta.StorageCopy:
 
         # Metadata should be checked for consistency before a job is accepted
@@ -491,9 +486,19 @@ class LoadDataFunc( _LoadSaveDataFunc, NodeFunction[_data.DataItem],):
 
     def _execute(self, ctx: NodeContext) -> _data.DataItem:
 
-        data_spec = _ctx_lookup(self.node.spec_id, ctx)
+        data_spec = self._choose_data_spec(self.node.spec_id, self.node.spec, ctx)
         data_copy = self._choose_copy(data_spec.data_item, data_spec.storage_def)
-        data_storage = self.storage.get_data_storage(data_copy.storageKey)
+
+        if data_spec.object_type == _api.ObjectType.DATA:
+            return self._load_data(data_spec, data_copy)
+
+        elif data_spec.object_type == _api.ObjectType.FILE:
+            return self._load_file(data_copy)
+
+        else:
+            raise _ex.EUnexpected()
+
+    def _load_data(self, data_spec, data_copy):
 
         trac_schema = data_spec.schema_def if data_spec.schema_def else data_spec.data_def.schema
         arrow_schema = _data.DataMapping.trac_to_arrow_schema(trac_schema) if trac_schema else None
@@ -503,36 +508,52 @@ class LoadDataFunc( _LoadSaveDataFunc, NodeFunction[_data.DataItem],):
         for opt_key, opt_value in data_spec.storage_def.storageOptions.items():
             options[opt_key] = _types.MetadataCodec.decode_value(opt_value)
 
-        table = data_storage.read_table(
+        storage = self.storage.get_data_storage(data_copy.storageKey)
+        table = storage.read_table(
             data_copy.storagePath,
             data_copy.storageFormat,
             arrow_schema,
             storage_options=options)
 
-        return _data.DataItem(table.schema, table)
+        return _data.DataItem(_api.ObjectType.DATA, table.schema, table)
 
+    def _load_file(self, data_copy):
 
-class SaveDataFunc(_LoadSaveDataFunc, NodeFunction[None]):
+        storage = self.storage.get_file_storage(data_copy.storageKey)
+        raw_bytes = storage.read_bytes(data_copy.storagePath)
+
+        return _data.DataItem(_api.ObjectType.FILE, raw_bytes=raw_bytes)
+
+
+class SaveDataFunc(_LoadSaveDataFunc, NodeFunction[_data.DataSpec]):
 
     def __init__(self, node: SaveDataNode, storage: _storage.StorageManager):
         super().__init__(storage)
         self.node = node
 
-    def _execute(self, ctx: NodeContext):
+    def _execute(self, ctx: NodeContext) -> _data.DataSpec:
 
         # Item to be saved should exist in the current context
         data_item = _ctx_lookup(self.node.data_item_id, ctx)
 
+        # Metadata already exists as data_spec but may not contain schema, row count, file size etc.
+        data_spec = self._choose_data_spec(self.node.spec_id, self.node.spec, ctx)
+        data_copy = self._choose_copy(data_spec.data_item, data_spec.storage_def)
+
         # Do not save empty outputs (optional outputs that were not produced)
         if data_item.is_empty():
-            return
+            return _data.DataSpec.create_empty_spec(data_item.object_type)
 
-        # This function assumes that metadata has already been generated as the data_spec
-        # i.e. it is already known which incarnation / copy of the data will be created
+        if data_item.object_type == _api.ObjectType.DATA:
+            return self._save_data(data_item, data_spec, data_copy)
 
-        data_spec = _ctx_lookup(self.node.spec_id, ctx)
-        data_copy = self._choose_copy(data_spec.data_item, data_spec.storage_def)
-        data_storage = self.storage.get_data_storage(data_copy.storageKey)
+        elif data_item.object_type == _api.ObjectType.FILE:
+            return self._save_file(data_item, data_spec, data_copy)
+
+        else:
+            raise _ex.EUnexpected()
+
+    def _save_data(self, data_item, data_spec, data_copy):
 
         # Current implementation will always put an Arrow table in the data item
         # Empty tables are allowed, so explicitly check if table is None
@@ -546,11 +567,32 @@ class SaveDataFunc(_LoadSaveDataFunc, NodeFunction[None]):
         for opt_key, opt_value in data_spec.storage_def.storageOptions.items():
             options[opt_key] = _types.MetadataCodec.decode_value(opt_value)
 
-        data_storage.write_table(
+        storage = self.storage.get_data_storage(data_copy.storageKey)
+        storage.write_table(
             data_copy.storagePath, data_copy.storageFormat,
             data_item.table,
             storage_options=options, overwrite=False)
 
+        data_spec = copy.deepcopy(data_spec)
+        # TODO: Save row count in metadata
+
+        if data_spec.data_def.schema is None and data_spec.data_def.schemaId is None:
+            data_spec.data_def.schema = _data.DataMapping.arrow_to_trac_schema(data_item.table.schema)
+
+        return data_spec
+
+    def _save_file(self, data_item, data_spec, data_copy):
+
+        if data_item.raw_bytes is None:
+            raise _ex.EUnexpected()
+
+        storage = self.storage.get_file_storage(data_copy.storageKey)
+        storage.write_bytes(data_copy.storagePath, data_item.raw_bytes)
+
+        data_spec = copy.deepcopy(data_spec)
+        data_spec.file_def.size = len(data_item.raw_bytes)
+
+        return data_spec
 
 def _model_def_for_import(import_details: meta.ImportModelJob):
 
@@ -571,8 +613,6 @@ class ImportModelFunc(NodeFunction[meta.ObjectDefinition]):
         self.node = node
         self._models = models
 
-        self._log = _util.logger_for_object(self)
-
     def _execute(self, ctx: NodeContext) -> meta.ObjectDefinition:
 
         model_stub = _model_def_for_import(self.node.import_details)
@@ -589,13 +629,15 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
             self, node: RunModelNode,
             model_class: _api.TracModel.__class__,
             checkout_directory: pathlib.Path,
-            storage_manager: _storage.StorageManager):
+            storage_manager: _storage.StorageManager,
+            log_provider: _logging.LogProvider):
 
         super().__init__()
         self.node = node
         self.model_class = model_class
         self.checkout_directory = checkout_directory
         self.storage_manager = storage_manager
+        self.log_provider = log_provider
 
     def _execute(self, ctx: NodeContext) -> Bundle[_data.DataView]:
 
@@ -622,7 +664,7 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
             for storage_key in self.node.storage_access:
                 if self.storage_manager.has_file_storage(storage_key, external=True):
                     storage_impl = self.storage_manager.get_file_storage(storage_key, external=True)
-                    storage = _ctx.TracFileStorageImpl(storage_key, storage_impl, write_access, self.checkout_directory)
+                    storage = _ctx.TracFileStorageImpl(storage_key, storage_impl, write_access, self.checkout_directory, self.log_provider)
                     storage_map[storage_key] = storage
                 elif self.storage_manager.has_data_storage(storage_key, external=True):
                     storage_impl = self.storage_manager.get_data_storage(storage_key, external=True)
@@ -630,7 +672,7 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
                     if not isinstance(storage_impl, _storage.IDataStorageBase):
                         raise _ex.EStorageConfig(f"External storage for [{storage_key}] is using the legacy storage framework]")
                     converter = _data.DataConverter.noop()
-                    storage = _ctx.TracDataStorageImpl(storage_key, storage_impl, converter, write_access, self.checkout_directory)
+                    storage = _ctx.TracDataStorageImpl(storage_key, storage_impl, converter, write_access, self.checkout_directory, self.log_provider)
                     storage_map[storage_key] = storage
                 else:
                     raise _ex.EStorageConfig(f"External storage is not available: [{storage_key}]")
@@ -642,12 +684,12 @@ class RunModelFunc(NodeFunction[Bundle[_data.DataView]]):
             trac_ctx = _ctx.TracDataContextImpl(
                 self.node.model_def, self.model_class,
                 local_ctx, dynamic_outputs, storage_map,
-                self.checkout_directory)
+                self.checkout_directory, self.log_provider)
         else:
             trac_ctx = _ctx.TracContextImpl(
                 self.node.model_def, self.model_class,
                 local_ctx, dynamic_outputs,
-                self.checkout_directory)
+                self.checkout_directory, self.log_provider)
 
         try:
             model = self.model_class()
@@ -750,9 +792,10 @@ class FunctionResolver:
 
     __ResolveFunc = tp.Callable[['FunctionResolver', Node[_T]], NodeFunction[_T]]
 
-    def __init__(self, models: _models.ModelLoader, storage: _storage.StorageManager):
+    def __init__(self, models: _models.ModelLoader, storage: _storage.StorageManager, log_provider: _logging.LogProvider):
         self._models = models
         self._storage = storage
+        self._log_provider = log_provider
 
     def resolve_node(self, node: Node[_T]) -> NodeFunction[_T]:
 
@@ -788,7 +831,7 @@ class FunctionResolver:
         checkout_directory = self._models.model_load_checkout_directory(node.model_scope, node.model_def)
         storage_manager = self._storage if node.storage_access else None
 
-        return RunModelFunc(node, model_class, checkout_directory, storage_manager)
+        return RunModelFunc(node, model_class, checkout_directory, storage_manager, self._log_provider)
 
     __basic_node_mapping: tp.Dict[Node.__class__, NodeFunction.__class__] = {
 
@@ -799,7 +842,6 @@ class FunctionResolver:
         DataViewNode: DataViewFunc,
         DataItemNode: DataItemFunc,
         BuildJobResultNode: BuildJobResultFunc,
-        SaveJobResultNode: SaveJobResultFunc,
         DataResultNode: DataResultFunc,
         StaticValueNode: StaticValueFunc,
         RuntimeOutputsNode: RuntimeOutputsFunc,

tracdap/rt/_exec/graph.py

@@ -13,7 +13,6 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-import pathlib
 import typing as tp
 import dataclasses as dc
 
@@ -182,15 +181,6 @@ class JobOutputs:
     bundles: tp.List[NodeId[ObjectBundle]] = dc.field(default_factory=list)
 
 
-# TODO: Where does this go?
-@dc.dataclass(frozen=True)
-class JobResultSpec:
-
-    save_result: bool = False
-    result_dir: tp.Union[str, pathlib.Path] = None
-    result_format: str = None
-
-
 # ----------------------------------------------------------------------------------------------------------------------
 #  NODE DEFINITIONS
 # ----------------------------------------------------------------------------------------------------------------------
@@ -309,20 +299,18 @@ class DataItemNode(MappingNode[_data.DataItem]):
 @_node_type
 class DataResultNode(Node[ObjectBundle]):
 
+    # TODO: Remove this node type
+    # Either produce metadata in SaveDataNode, or handle DataSpec outputs in result processing nodes
+
     output_name: str
-    data_item_id: NodeId[_data.DataItem]
-    data_spec_id: NodeId[_data.DataSpec]
-    data_save_id: NodeId[type(None)]
+    data_save_id: NodeId[_data.DataSpec]
 
-    data_key: str
-    storage_key: str
+    data_key: str = None
+    file_key: str = None
+    storage_key: str = None
 
     def _node_dependencies(self) -> tp.Dict[NodeId, DependencyType]:
-
-        return {
-            self.data_item_id: DependencyType.HARD,
-            self.data_spec_id: DependencyType.HARD,
-            self.data_save_id: DependencyType.HARD}
+        return {self.data_save_id: DependencyType.HARD}
 
 
 @_node_type
@@ -333,24 +321,33 @@ class LoadDataNode(Node[_data.DataItem]):
     The latest incarnation of the item will be loaded from any available copy
     """
 
-    spec_id: NodeId[_data.DataSpec]
+    spec_id: tp.Optional[NodeId[_data.DataSpec]] = None
+    spec: tp.Optional[_data.DataSpec] = None
 
     def _node_dependencies(self) -> tp.Dict[NodeId, DependencyType]:
-        return {self.spec_id: DependencyType.HARD}
+        deps = dict()
+        if self.spec_id is not None:
+            deps[self.spec_id] = DependencyType.HARD
+        return deps
 
 
 @_node_type
-class SaveDataNode(Node[None]):
+class SaveDataNode(Node[_data.DataSpec]):
 
     """
     Save an individual data item to storage
     """
 
-    spec_id: NodeId[_data.DataSpec]
     data_item_id: NodeId[_data.DataItem]
 
+    spec_id: tp.Optional[NodeId[_data.DataSpec]] = None
+    spec: tp.Optional[_data.DataSpec] = None
+
     def _node_dependencies(self) -> tp.Dict[NodeId, DependencyType]:
-        return {self.spec_id: DependencyType.HARD, self.data_item_id: DependencyType.HARD}
+        deps = {self.data_item_id: DependencyType.HARD}
+        if self.spec_id is not None:
+            deps[self.spec_id] = DependencyType.HARD
+        return deps
 
 
 @_node_type
@@ -395,6 +392,7 @@ class RuntimeOutputsNode(Node[JobOutputs]):
 @_node_type
 class BuildJobResultNode(Node[cfg.JobResult]):
 
+    result_id: meta.TagHeader
     job_id: meta.TagHeader
 
     outputs: JobOutputs
@@ -407,16 +405,6 @@ class BuildJobResultNode(Node[cfg.JobResult]):
         return {node_id: DependencyType.HARD for node_id in dep_ids}
 
 
-@_node_type
-class SaveJobResultNode(Node[None]):
-
-    job_result_id: NodeId[cfg.JobResult]
-    result_spec: JobResultSpec
-
-    def _node_dependencies(self) -> tp.Dict[NodeId, DependencyType]:
-        return {self.job_result_id: DependencyType.HARD}
-
-
 @_node_type
 class ChildJobNode(Node[cfg.JobResult]):