Skip to content

memory

Implementation of in-memory artifact store.

InMemoryArtifactMapper

Bases: ArtifactMapper

In-memory mapper for the artifact resource.

Source code in mlte/store/artifact/underlying/memory.py
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
class InMemoryArtifactMapper(ArtifactMapper):
    """In-memory mapper for the artifact resource."""

    def __init__(
        self, *, storage: MemoryArtifactStorage, validators: CompositeValidator
    ) -> None:
        super().__init__(validators=validators)

        self.storage = storage
        """A reference to underlying storage."""

    def read(
        self, artifact_id: str, model_and_version: tuple[str, str]
    ) -> ArtifactModel:
        model_id, version_id = model_and_version
        return self._get_artifact(model_id, version_id, artifact_id)

    def list(self, model_and_version: tuple[str, str]) -> list[str]:
        model_id, version_id = model_and_version
        return [
            artifact_id
            for artifact_id in self._get_artifacts(model_id, version_id)
        ]

    def delete(
        self, artifact_id: str, model_and_version: tuple[str, str]
    ) -> ArtifactModel:
        model_id, version_id = model_and_version
        artifact = self.read(artifact_id, (model_id, version_id))

        if artifact.header.level == ArtifactLevel.VERSION:
            self.storage.models[model_id].delete_artifact(
                artifact_id, version_id
            )
        else:
            self.storage.models[model_id].delete_artifact(artifact_id)

        return artifact

    # -------------------------------------------------------------------------
    # Internal helpers.
    # -------------------------------------------------------------------------

    def _store_artifact(
        self, artifact: ArtifactModel, model_and_version: tuple[str, str]
    ):
        """Writes an artifact to the store."""
        model_id, version_id = model_and_version
        self.storage.add_artifact(
            artifact, model_id, version_id, artifact.header.level
        )
        return self.read(artifact.header.identifier, (model_id, version_id))

    def _get_artifact(
        self, model_id: str, version_id: str, artifact_id: str
    ) -> ArtifactModel:
        """
        Get the artifact from storage.
        :param model_id: The identifier for the model.
        :param version_id: The identifier for the version.
        :param artifact_id: The identifier for the artifact.
        :raises ErrorNotFound: If the model or version are not present.
        :return: The artifcat.
        """
        if model_id not in self.storage.models:
            raise errors.ErrorNotFound(f"Model {model_id}")
        model_artifacts = self.storage.models[model_id]

        if version_id not in model_artifacts.version_artifacts:
            raise errors.ErrorNotFound(
                f"Version: {version_id} for model {model_id}"
            )

        artifact = model_artifacts.get_artifact(version_id, artifact_id)
        if not artifact:
            raise errors.ErrorNotFound(
                f"Artifact: {artifact_id} for version {version_id} and model {model_id}"
            )
        else:
            return artifact

    def _get_artifacts(
        self, model_id: str, version_id: str
    ) -> OrderedDict[str, ArtifactModel]:
        """
        Get the artifcats from storage.
        :param model_id: The identifier for the model.
        :param version_id: The identifier for the version.
        :raises ErrorNotFound: If the model or version are not present.
        :return: The artifcats for this model and version, including model-level only.
        """
        if model_id not in self.storage.models:
            raise errors.ErrorNotFound(f"Model {model_id}")
        model_artifacts = self.storage.models[model_id]

        if version_id not in model_artifacts.version_artifacts:
            raise errors.ErrorNotFound(f"Version {version_id}")

        return model_artifacts.get_all_artifacts(version_id)

storage = storage instance-attribute

A reference to underlying storage.

InMemoryModelMapper

Bases: ModelMapper

In-memory mapper for the model resource.

Source code in mlte/store/artifact/underlying/memory.py
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
class InMemoryModelMapper(ModelMapper):
    """In-memory mapper for the model resource."""

    def __init__(
        self,
        *,
        storage: MemoryArtifactStorage,
        version_mapper: InMemoryVersionMapper,
    ) -> None:
        self.storage = storage
        """A reference to underlying storage."""

        self.version_mapper = version_mapper
        """The mapper to version CRUD."""

    def create(self, model: Model, context: Any = None) -> Model:
        if model.identifier in self.storage.models:
            raise errors.ErrorAlreadyExists(f"Model {model.identifier}")
        self.storage.models[model.identifier] = ModelArtifacts()
        return Model(identifier=model.identifier, versions=[])

    def read(self, model_id: str, context: Any = None) -> Model:
        if model_id not in self.storage.models:
            raise errors.ErrorNotFound(f"Model {model_id}")
        return Model(
            identifier=model_id,
            versions=[
                self.version_mapper.read(id, model_id)
                for id in self.storage.models[model_id].version_artifacts.keys()
            ],
        )

    def list(self, context: Any = None) -> list[str]:
        return [model_id for model_id in self.storage.models.keys()]

    def delete(self, model_id: str, context: Any = None) -> Model:
        popped = self.read(model_id)
        del self.storage.models[model_id]
        return popped

storage = storage instance-attribute

A reference to underlying storage.

version_mapper = version_mapper instance-attribute

The mapper to version CRUD.

InMemoryStore

Bases: ArtifactStore

An in-memory implementation of the MLTE artifact store.

Source code in mlte/store/artifact/underlying/memory.py
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
class InMemoryStore(ArtifactStore):
    """An in-memory implementation of the MLTE artifact store."""

    def __init__(self, uri: StoreURI) -> None:
        super().__init__(uri=uri)

        self.storage = MemoryArtifactStorage()
        """The underlying storage for the store."""

    def session(self) -> InMemoryStoreSession:
        """
        Return a session handle for the store instance.
        :return: The session handle
        """
        return InMemoryStoreSession(
            storage=self.storage, validators=self.validators
        )

storage = MemoryArtifactStorage() instance-attribute

The underlying storage for the store.

session()

Return a session handle for the store instance.

Returns:

Type Description
InMemoryStoreSession

The session handle

Source code in mlte/store/artifact/underlying/memory.py
106
107
108
109
110
111
112
113
def session(self) -> InMemoryStoreSession:
    """
    Return a session handle for the store instance.
    :return: The session handle
    """
    return InMemoryStoreSession(
        storage=self.storage, validators=self.validators
    )

InMemoryStoreSession

Bases: ArtifactStoreSession

An in-memory implementation of the MLTE artifact store.

Source code in mlte/store/artifact/underlying/memory.py
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
class InMemoryStoreSession(ArtifactStoreSession):
    """An in-memory implementation of the MLTE artifact store."""

    def __init__(
        self, *, storage: MemoryArtifactStorage, validators: CompositeValidator
    ) -> None:
        self.storage = storage
        """A reference to underlying storage."""

        self.version_mapper = InMemoryVersionMapper(storage=storage)
        """The mapper to version CRUD."""

        self.model_mapper = InMemoryModelMapper(
            storage=storage, version_mapper=self.version_mapper
        )
        """The mapper to model CRUD."""

        self.artifact_mapper = InMemoryArtifactMapper(
            storage=storage, validators=validators
        )
        """The mapper to artifact CRUD."""

    def close(self) -> None:
        """Close the session."""
        # Closing an in-memory session is a no-op.
        pass

artifact_mapper = InMemoryArtifactMapper(storage=storage, validators=validators) instance-attribute

The mapper to artifact CRUD.

model_mapper = InMemoryModelMapper(storage=storage, version_mapper=(self.version_mapper)) instance-attribute

The mapper to model CRUD.

storage = storage instance-attribute

A reference to underlying storage.

version_mapper = InMemoryVersionMapper(storage=storage) instance-attribute

The mapper to version CRUD.

close()

Close the session.

Source code in mlte/store/artifact/underlying/memory.py
138
139
140
141
def close(self) -> None:
    """Close the session."""
    # Closing an in-memory session is a no-op.
    pass

InMemoryVersionMapper

Bases: VersionMapper

In-memory mapper for the version resource.

Source code in mlte/store/artifact/underlying/memory.py
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
class InMemoryVersionMapper(VersionMapper):
    """In-memory mapper for the version resource."""

    def __init__(self, *, storage: MemoryArtifactStorage) -> None:
        self.storage = storage
        """A reference to underlying storage."""

    def create(self, version: Version, model_id: str) -> Version:
        if model_id not in self.storage.models:
            raise errors.ErrorNotFound(f"Model {model_id}")

        model = self.storage.models[model_id]
        if version.identifier in model.version_artifacts:
            raise errors.ErrorAlreadyExists(f"Version {version.identifier}")

        model.version_artifacts[version.identifier] = OrderedDict()
        return Version(identifier=version.identifier)

    def read(self, version_id: str, model_id: str) -> Version:
        if model_id not in self.storage.models:
            raise errors.ErrorNotFound(f"Model {model_id}")

        model = self.storage.models[model_id]
        if version_id not in model.version_artifacts:
            raise errors.ErrorNotFound(
                f"Version {version_id} for model {model_id}"
            )

        return Version(identifier=version_id)

    def list(self, model_id: str) -> list[str]:
        if model_id not in self.storage.models:
            raise errors.ErrorNotFound(f"Model {model_id}")

        model = self.storage.models[model_id]
        return [version_id for version_id in model.version_artifacts.keys()]

    def delete(self, version_id: str, model_id: str) -> Version:
        popped = self.read(version_id, model_id)
        model = self.storage.models[model_id]
        del model.version_artifacts[version_id]
        return popped

storage = storage instance-attribute

A reference to underlying storage.

MemoryArtifactStorage

A simple storage wrapper for the in-memory store.

Source code in mlte/store/artifact/underlying/memory.py
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
class MemoryArtifactStorage:
    """A simple storage wrapper for the in-memory store."""

    def __init__(self) -> None:
        self.models: dict[str, ModelArtifacts] = {}

    def add_artifact(
        self,
        artifact: ArtifactModel,
        model_id: str,
        version_id: str,
        level: ArtifactLevel,
    ):
        """Adds an artifact to the model or version level list."""
        if level == ArtifactLevel.MODEL:
            self.models[model_id].artifacts[
                artifact.header.identifier
            ] = artifact
        else:
            self.models[model_id].version_artifacts[version_id][
                artifact.header.identifier
            ] = artifact

add_artifact(artifact, model_id, version_id, level)

Adds an artifact to the model or version level list.

Source code in mlte/store/artifact/underlying/memory.py
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def add_artifact(
    self,
    artifact: ArtifactModel,
    model_id: str,
    version_id: str,
    level: ArtifactLevel,
):
    """Adds an artifact to the model or version level list."""
    if level == ArtifactLevel.MODEL:
        self.models[model_id].artifacts[
            artifact.header.identifier
        ] = artifact
    else:
        self.models[model_id].version_artifacts[version_id][
            artifact.header.identifier
        ] = artifact

ModelArtifacts

A structure that contains model artifacts, at top level and version level.

Source code in mlte/store/artifact/underlying/memory.py
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class ModelArtifacts:
    """A structure that contains model artifacts, at top level and version level."""

    def __init__(self) -> None:
        self.artifacts: OrderedDict[str, ArtifactModel] = OrderedDict()
        """The artifacts associated with the model only."""

        self.version_artifacts: dict[str, OrderedDict[str, ArtifactModel]] = {}
        """The version level artifacts, by version."""

    def get_all_artifacts(
        self, version_id: str
    ) -> OrderedDict[str, ArtifactModel]:
        """Returns all artifacts, from model and version level."""
        all_artifacts = self.artifacts.copy()
        all_artifacts.update(self.version_artifacts[version_id])
        return all_artifacts

    def get_artifact(
        self, version_id: str, artifact_id: str
    ) -> Optional[ArtifactModel]:
        """Gets an artifact from a model or version level."""
        # First check at the model level.
        if artifact_id in self.artifacts:
            return self.artifacts[artifact_id]
        else:
            # Then check at the version level.
            if artifact_id in self.version_artifacts[version_id]:
                return self.version_artifacts[version_id][artifact_id]
            else:
                return None

    def delete_artifact(
        self, artifact_id: str, version_id: Optional[str] = None
    ):
        """Removes the given artifact, from the given version, or from the model list."""
        if version_id:
            del self.version_artifacts[version_id][artifact_id]
        else:
            del self.artifacts[artifact_id]

artifacts = OrderedDict() instance-attribute

The artifacts associated with the model only.

version_artifacts = {} instance-attribute

The version level artifacts, by version.

delete_artifact(artifact_id, version_id=None)

Removes the given artifact, from the given version, or from the model list.

Source code in mlte/store/artifact/underlying/memory.py
58
59
60
61
62
63
64
65
def delete_artifact(
    self, artifact_id: str, version_id: Optional[str] = None
):
    """Removes the given artifact, from the given version, or from the model list."""
    if version_id:
        del self.version_artifacts[version_id][artifact_id]
    else:
        del self.artifacts[artifact_id]

get_all_artifacts(version_id)

Returns all artifacts, from model and version level.

Source code in mlte/store/artifact/underlying/memory.py
36
37
38
39
40
41
42
def get_all_artifacts(
    self, version_id: str
) -> OrderedDict[str, ArtifactModel]:
    """Returns all artifacts, from model and version level."""
    all_artifacts = self.artifacts.copy()
    all_artifacts.update(self.version_artifacts[version_id])
    return all_artifacts

get_artifact(version_id, artifact_id)

Gets an artifact from a model or version level.

Source code in mlte/store/artifact/underlying/memory.py
44
45
46
47
48
49
50
51
52
53
54
55
56
def get_artifact(
    self, version_id: str, artifact_id: str
) -> Optional[ArtifactModel]:
    """Gets an artifact from a model or version level."""
    # First check at the model level.
    if artifact_id in self.artifacts:
        return self.artifacts[artifact_id]
    else:
        # Then check at the version level.
        if artifact_id in self.version_artifacts[version_id]:
            return self.version_artifacts[version_id][artifact_id]
        else:
            return None