Skip to content

Parameter store

StepRequest dataclass

Declarative request for creating/updating datasets by parameter set.

Source code in kosh/parameter_store.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
@dataclass(frozen=True)
class StepRequest:
    """Declarative request for creating/updating datasets by parameter set."""

    store_uri: str
    step: str
    ensemble_name: Optional[str] = None
    init_params: Mapping[str, Any] = field(default_factory=dict)
    dataset_record_type: Optional[str] = None
    rtol: float = 1e-5
    atol: float = 1e-8
    strict_match: bool = False
    upsert_init: bool = False
    delete_ensemble: bool = False
    wipe_store: bool = False
    step_field: str = "workflow_step"
    init_step_value: str = "init"
    metadata_updates: Mapping[str, Any] = field(default_factory=dict)
    dataset_associations: Sequence[Association] = field(default_factory=tuple)
    ensemble_associations: Sequence[Association] = field(default_factory=tuple)
    ensemble_metadata_updates: Mapping[str, Any] = field(default_factory=dict)
    update_record: Optional[Mapping[str, Any]] = None
    connect_kwargs: Mapping[str, Any] = field(default_factory=dict)

apply_step(request)

Create or update dataset(s) in a Kosh store for a workflow step.

Behavior: - If request.step == request.init_step_value: create a dataset with init_params (and add to the ensemble when request.ensemble_name is provided). - Otherwise: find dataset(s) matching init_params (DataRange for numerics), then update. If request.ensemble_name is provided, restrict the search to that ensemble.

Source code in kosh/parameter_store.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def apply_step(request: StepRequest) -> List[str]:
    """Create or update dataset(s) in a Kosh store for a workflow step.

    Behavior:
    - If `request.step == request.init_step_value`: create a dataset with `init_params` (and add to the
      ensemble when `request.ensemble_name` is provided).
    - Otherwise: find dataset(s) matching `init_params` (DataRange for numerics), then update. If
      `request.ensemble_name` is provided, restrict the search to that ensemble.
    """
    store = _open_store(request.store_uri, wipe_store=request.wipe_store, connect_kwargs=request.connect_kwargs)
    try:
        record_type = request.dataset_record_type
        find_kwargs: Dict[str, Any] = {}
        if record_type:
            find_kwargs["types"] = [record_type]

        ensemble = None
        if request.ensemble_name is None:
            if request.delete_ensemble:
                raise ValueError("delete_ensemble requires an ensemble_name")
            if request.ensemble_associations:
                raise ValueError("ensemble_associations requires an ensemble_name")
            if request.ensemble_metadata_updates:
                raise ValueError("ensemble_metadata_updates requires an ensemble_name")
        else:
            ensemble = _get_or_create_ensemble(
                store, name=request.ensemble_name, delete_ensemble=request.delete_ensemble
            )

            for assoc in request.ensemble_associations:
                ensemble.associate(assoc.path, assoc.mime_type)
            if request.ensemble_metadata_updates:
                ensemble.update(dict(request.ensemble_metadata_updates))

        if request.step == request.init_step_value:
            dataset_ids: List[str] = []
            if request.upsert_init:
                query = _build_query(
                    request.init_params,
                    rtol=request.rtol,
                    atol=request.atol,
                    strict_match=request.strict_match,
                )
                if ensemble is None:
                    dataset_ids = list(store.find(ids_only=True, **find_kwargs, **query))
                else:
                    dataset_ids = list(ensemble.find_datasets(ids_only=True, **query))
                if len(dataset_ids) > 1:
                    raise ValueError(f"Found more than one dataset {len(dataset_ids)} matching init params, giving up")
            if not dataset_ids:
                metadata = dict(request.init_params)
                metadata[request.step_field] = request.init_step_value
                dataset = store.create(metadata=metadata)
                if ensemble is not None:
                    ensemble.add(dataset)
                dataset_ids = [dataset.id]
        else:
            query = _build_query(
                request.init_params,
                rtol=request.rtol,
                atol=request.atol,
                strict_match=request.strict_match,
            )
            if ensemble is None:
                dataset_ids = list(store.find(ids_only=True, **find_kwargs, **query))
            else:
                dataset_ids = list(ensemble.find_datasets(ids_only=True, **query))

        if not dataset_ids:
            raise ValueError("could not find any dataset that match your request")

        for dataset_id in dataset_ids:
            ds = store.open(dataset_id)
            meta = dict(request.metadata_updates)
            meta[request.step_field] = request.step
            if request.update_record:
                meta.update(dict(request.update_record))
            if meta:
                ds.update(meta)
            for assoc in request.dataset_associations:
                ds.associate(assoc.path, assoc.mime_type)

        return dataset_ids
    finally:
        store.close()