14. Export Models

This tutorial is based on example code which can be found in the finTRAC Examples Repository.

The placements in Exporting Data write datasets to external storage exactly as they are. When an export needs custom logic, such as reshaping data for a downstream system, adding fields or writing a non-standard format, it can be done with an export model. An export model runs as a node in a flow. It receives datasets from the flow like any other model, and it has write access to the external storage locations the job grants it.

This tutorial adds an export model to the flow from Chaining Models.

Note

Export models are part of the experimental API in tracdap.rt.api.experimental, which may change in future versions of TRAC.

Writing an export model

Export models extend TracDataExport instead of TracModel. Parameters and inputs are defined in the usual way. Export models do not have outputs, everything they produce is written to external storage.

src/tutorial/data_export.py
16import tracdap.rt.api.experimental as trac
17
18import tutorial.schemas as schemas
19
20
21class DataExportExample(trac.TracDataExport):
22
23    def define_parameters(self):
24
25        return trac.define_parameters(
26            trac.P("storage_key", trac.STRING, "TRAC external storage key"),
27            trac.P("no_overwrite", trac.BOOLEAN, "Do not overwrite any files that already exist", default_value=False),
28            trac.P("export_comment", trac.STRING, "Comment for this export operation (optional)", default_value=""))
29
30    def define_inputs(self):
31
32        profit_by_region = trac.load_schema(schemas, "profit_by_region.csv")
33
34        return {"profit_by_region": trac.define_input(profit_by_region, label="Profit by region")}

The model is given the storage key to write to as a parameter, storage_key. In run_model(), the context object is a TracDataContext, which has a get_file_storage() method to access a storage location by its key. The TracFileStorage interface it returns can create directories, check for existing files and write files as streams:

36    def run_model(self, ctx: trac.TracDataContext):
37
38        # Get a dataset - set up as inputs the same as model run jobs
39        profit_by_region = ctx.get_pandas_table("profit_by_region")
40
41        # Add some fields to label the export
42        pbr_augmented = self.add_export_metadata("profit_by_region", profit_by_region, ctx)
43
44        # Get access to the storage interface - must be in the TRAC config as a storage location
45        storage_key = ctx.get_parameter("storage_key")
46        storage = ctx.get_file_storage(storage_key)
47
48        # Storage path, relative to the root of the storage location
49        export_dir = "data_export_example"
50        export_file = export_dir + "/profit_by_region.csv"
51
52        storage.mkdir(export_dir, recursive=True)
53
54        no_overwrite = ctx.get_parameter("no_overwrite")
55
56        if no_overwrite and storage.exists(export_file):
57            ctx.log.info(f"Export of [{export_file}] will be skipped because the file already exists")
58            return
59
60        with storage.write_byte_stream(export_file) as stream:
61            pbr_augmented.to_csv(stream, index=False)
62
63    @staticmethod
64    def add_export_metadata(dataset_name, dataset, ctx: trac.TracDataContext):
65
66        export_comment = ctx.get_parameter("export_comment")
67
68        return dataset.assign(
69            dataset_name=dataset_name,
70            comments=export_comment)

Adding the export to a flow

The flow is the same as the one in Chaining Models, with one extra node for the export model. The node sets modelType to DATA_EXPORT_MODEL, and its input is connected to the profit_by_region output of the aggregation model:

config/chaining_export_flow.yaml
 2nodes:
 3
 4  customer_loans:
 5    nodeType: "INPUT_NODE"
 6
 7  customer_data_filter:
 8    nodeType: "MODEL_NODE"
 9    inputs: [customer_loans]
10    outputs: [filtered_loans]
11
12  pnl_aggregation:
13    nodeType: "MODEL_NODE"
14    inputs: [customer_loans]
15    outputs: [profit_by_region]
16
17  data_export:
18    nodeType: "MODEL_NODE"
19    modelType: "DATA_EXPORT_MODEL"
20    inputs: [profit_by_region]
21
22  profit_by_region:
23    nodeType: "OUTPUT_NODE"
24
25
26edges:
27
28  - source: { node: customer_data_filter, socket: filtered_loans }
29    target: { node: pnl_aggregation, socket: customer_loans }

The job config adds the export model and its parameters to the chaining job. The exportStorageAccess list says which external storage locations export models in the flow are allowed to write to. A storage key passed to an export model has to be in this list, or the model will not be able to access it.

config/chaining_export.yaml
 2job:
 3  runFlow:
 4
 5    flow: ./chaining_export_flow.yaml
 6
 7    parameters:
 8      eur_usd_rate: 1.2071
 9      default_weighting: 1.5
10      filter_defaults: false
11      filter_region: munster
12      storage_key: exported_data
13      export_comment: "Profit by region, exported from the chaining flow"
14
15    inputs:
16      customer_loans: "inputs/loan_final313_100.csv"
17
18    outputs:
19      profit_by_region: "outputs/chaining_export/profit_by_region.csv"
20
21    models:
22      customer_data_filter: tutorial.chaining.CustomerDataFilter
23      pnl_aggregation: tutorial.using_data.PnlAggregation
24      data_export: tutorial.data_export.DataExportExample
25
26    exportStorageAccess:
27      - exported_data

The flow can be launched in the same way as in Chaining Models:

73if __name__ == "__main__":
74    import tracdap.rt.launch as launch
75    launch.launch_job("config/chaining_export.yaml", "config/sys_config.yaml", dev_mode=True)

See also

Full source code is available for the Export model example on GitHub.