def _get_artifact(self, model_id: str) -> ArtifactCollection:
if not self._bucket_exists():
raise NoSuchArtifactError(model_id, self)
keys = list(self._list_blobs(model_id).keys())
if len(keys) == 0:
raise NoSuchArtifactError(model_id, self)
elif len(keys) == 1 and keys[0] == '.':
return Blobs({})
else:
return Blobs({
os.path.relpath(key, model_id): S3Blob(key, self.bucket_name, self.endpoint)
for key in keys
})
def dump(self, model) -> ArtifactCollection:
"""
Dumps model artifacts as :class:`~ebonite.core.objects.ArtifactCollection`
:return: context manager with :class:`~ebonite.core.objects.ArtifactCollection`
"""
model_blob, refs = self._serialize_model(model)
blobs = {self.model_filename: InMemoryBlob(model_blob)}
artifact_cms = []
uuids = []
for uuid, (io, obj) in refs.items():
blobs[uuid + self.io_ext] = InMemoryBlob(self._serialize_io(io))
artifact_cms.append(io.dump(obj))
uuids.append(uuid)
from ebonite.core.objects.artifacts import _enter_all_cm, _ExitAllCm, _RelativePathWrapper
additional_artifacts = _enter_all_cm(artifact_cms)
with _ExitAllCm(artifact_cms):
additional_artifacts = [
_RelativePathWrapper(art, uuid)
for art, uuid in zip(additional_artifacts, uuids)
]
yield CompositeArtifactCollection([Blobs(blobs)] +
additional_artifacts)
def dump(self, session, path) -> FilesContextManager:
# path = os.path.join(path, TF_MODEL_FILENAME) # TODO check if it is needed
saver = tf.train.Saver()
saver.save(session, path)
yield Blobs({
name: LocalFileBlob(os.path.join(path, name))
for name in os.listdir(path)
})
def dump(self) -> FilesContextManager:
with self.io.dump(self.model) as artifact:
yield artifact + Blobs({
self.methods_json:
InMemoryBlob(dumps(self.methods).encode('utf-8')),
self.requirements_json:
InMemoryBlob(dumps(self.requirements).encode('utf-8'))
})
def _push_artifact(self, model_id: str,
blobs: typing.Dict[str, Blob]) -> ArtifactCollection:
if model_id in self._cache:
raise ArtifactExistsError(model_id, self)
self._cache[model_id] = Blobs(
{k: InMemoryBlob(v.bytes())
for k, v in blobs.items()})
return self._cache[model_id]
def dump(self, session, path) -> FilesContextManager:
with session.as_default(), session.graph.as_default():
saver = tf.train.Saver(save_relative_paths=True)
saver.save(session, os.path.join(path, TF_MODEL_FILENAME))
yield Blobs({
name: LocalFileBlob(os.path.join(path, name)) for name in os.listdir(path)
})
def _get_artifact(self, model_id: str) -> ArtifactCollection:
path = os.path.join(self.path, model_id)
if not os.path.exists(path):
raise NoSuchArtifactError(model_id, self)
return Blobs({
os.path.relpath(file, path): LocalFileBlob(os.path.join(self.path, file)) for file in
glob.glob(os.path.join(path, '**'), recursive=True) if os.path.isfile(file)
})
def dump(self) -> ArtifactCollection:
"""
Dumps `torch.nn.Module` instance to :class:`.InMemoryBlob` and creates :class:`.ArtifactCollection` from it
:return: context manager with :class:`~ebonite.core.objects.ArtifactCollection`
"""
buffer = BytesIO()
torch.save(self.model, buffer)
yield Blobs({self.model_file_name: InMemoryBlob(buffer.getvalue())})
def dump(self, session, path) -> FilesContextManager:
tf.train.write_graph(session.graph.as_graph_def(),
path,
TF_MODEL_FILENAME,
as_text=False)
yield Blobs({
TF_MODEL_FILENAME:
LocalFileBlob(os.path.join(path, TF_MODEL_FILENAME))
})
def push_artifact(self, artifact_type, artifact_id: str,
blobs: typing.Dict[str, Blob]) -> ArtifactCollection:
artifact_id = f'{artifact_type}/{artifact_id}'
if artifact_id in self._cache:
raise ArtifactExistsError(artifact_id, self)
self._cache[artifact_id] = Blobs(
{k: InMemoryBlob(v.bytes())
for k, v in blobs.items()})
return self._cache[artifact_id]
def _push_artifact(self, model_id: str, blobs: typing.Dict[str, Blob]) -> ArtifactCollection:
self._ensure_bucket()
if len(self._list_blobs(model_id)) > 0:
raise ArtifactExistsError(model_id, self)
if len(blobs) == 0:
self._s3.upload_fileobj(io.BytesIO(b''), self.bucket_name, model_id)
return Blobs({})
result = {}
for filepath, blob in blobs.items():
join = os.path.join(model_id, filepath)
with blob.bytestream() as b:
logger.debug('Uploading %s to s3 %s/%s', blob, self.endpoint, self.bucket_name)
self._s3.upload_fileobj(b, self.bucket_name, join)
result[filepath] = S3Blob(join, self.bucket_name, self.endpoint)
return Blobs(result)
def dump(self, model: _TfModel) -> FilesContextManager:
"""
Dumps session to temporary directory and creates :class:`~ebonite.core.objects.ArtifactCollection` from it
:return: context manager with :class:`~ebonite.core.objects.ArtifactCollection`
"""
with tempfile.TemporaryDirectory(prefix='ebonite_tensor_') as tempdir:
dumper = self._get_dumper(model.is_frozen)
with dumper.dump(model.get_session(), tempdir) as artifact:
meta = json.dumps([model.tensor_names, model.is_frozen]).encode('utf-8')
yield artifact + Blobs({self.meta_json: InMemoryBlob(meta)})
def dump(self, model) -> FilesContextManager:
with tempfile.TemporaryDirectory(prefix='ebonite_tf_v2') as tmpdir:
dir_path = os.path.join(tmpdir, self.model_dir_name)
model.save(dir_path)
shutil.make_archive(dir_path, 'zip', dir_path)
yield Blobs({
self.model_dir_name + self.ext:
LocalFileBlob(dir_path + self.ext)
})
def get_artifacts(self) -> ArtifactCollection:
"""Return model binaries"""
artifacts = CompositeArtifactCollection(
[_RelativePathWrapper(m.artifact_any, os.path.join(MODEL_BIN_PATH, m.name)) for m in
self.pipeline.models.values()])
if len(self.server.additional_binaries) > 0:
artifacts = CompositeArtifactCollection([
artifacts,
Blobs({os.path.basename(f): LocalFileBlob(f) for f in self.server.additional_binaries})
])
return artifacts
def test_task__delete_model_with_artifacts(task_saved, model, artifact_repo):
model._unpersisted_artifacts = Blobs({'data': InMemoryBlob(b'data')})
task_saved.bind_artifact_repo(artifact_repo)
task_saved.push_model(model)
assert model.id is not None
assert model.task_id is not None
task_saved.delete_model(model)
assert len(task_saved.models) == 0
assert model.id is None
assert model.task_id is None
def dump(self, model) -> ArtifactCollection:
"""
Dumps `catboost.CatBoostClassifier` or `catboost.CatBoostRegressor` instance to :class:`.LocalFileBlob` and
creates :class:`.ArtifactCollection` from it
:return: context manager with :class:`~ebonite.core.objects.ArtifactCollection`
"""
model_file = tempfile.mktemp()
try:
model.save_model(model_file)
yield Blobs({self._get_model_file_name(model): LocalFileBlob(model_file)})
finally:
os.remove(model_file)
def dump(self, model) -> ArtifactCollection:
"""
Dumps `torch.nn.Module` instance to :class:`.InMemoryBlob` and creates :class:`.ArtifactCollection` from it
:return: context manager with :class:`~ebonite.core.objects.ArtifactCollection`
"""
is_jit = isinstance(model, torch.jit.ScriptModule)
save = torch.jit.save if is_jit else torch.save
model_name = self.model_jit_file_name if is_jit else self.model_file_name
buffer = BytesIO()
save(model, buffer)
yield Blobs({model_name: InMemoryBlob(buffer.getvalue())})
def artifact_collection():
blobs = Blobs({
'1': InMemoryBlob(bytes(123)),
'2': InMemoryBlob(bytes(321))
})
ac = blobs + _RelativePathWrapper(blobs, 'first') + _RelativePathWrapper(
blobs, 'second')
ac += _RelativePathWrapper(blobs, 'third')
ret = ac + _RelativePathWrapper(ac, 'go')
ret += _RelativePathWrapper(ac, 'be')
return ret
def get_artifacts(self) -> ArtifactCollection:
"""Return model binaries"""
artifacts = _RelativePathWrapper(self.model.artifact_any,
MODEL_BIN_PATH)
if len(self.server.additional_binaries) > 0:
artifacts = CompositeArtifactCollection([
artifacts,
Blobs({
os.path.basename(f): LocalFileBlob(f)
for f in self.server.additional_binaries
})
])
return artifacts
def _push_artifact(self, model_id: str, blobs: typing.Dict[str, Blob]) -> ArtifactCollection:
path = os.path.join(self.path, model_id)
if os.path.exists(path):
raise ArtifactExistsError(model_id, self)
os.makedirs(path, exist_ok=True)
result = {}
for filepath, blob in blobs.items():
join = os.path.join(path, filepath)
os.makedirs(os.path.dirname(join), exist_ok=True)
logger.debug('Writing artifact %s to %s', blob, join)
blob.materialize(join)
result[filepath] = LocalFileBlob(join)
return Blobs(result)
def dump(self) -> ArtifactCollection:
"""
Dumps model artifacts as :class:`~ebonite.core.objects.ArtifactCollection`
:return: context manager with :class:`~ebonite.core.objects.ArtifactCollection`
"""
model_blob, refs = self._serialize_model()
blobs = {self.model_filename: InMemoryBlob(model_blob)}
additional_artifacts = []
for uuid, wrapper in refs.items():
blobs[uuid + self.wrapper_ext] = InMemoryBlob(
self._serialize_wrapper(wrapper))
with wrapper.dump() as artifact:
additional_artifacts.append(
_RelativePathWrapper(artifact, uuid))
yield CompositeArtifactCollection([Blobs(blobs)] +
additional_artifacts)
def dump(self, model) -> FilesContextManager:
yield Blobs({'test.bin': InMemoryBlob(b'test')})
def write(self,
dataset: Dataset) -> Tuple[DatasetReader, ArtifactCollection]:
return OneFileDatasetReader(dataset.dataset_type), \
Blobs.from_blobs({self.FILENAME: InMemoryBlob(self.convert(dataset.data))})
def get_artifacts(self):
return Blobs({'test.bin': InMemoryBlob(b'test_bytes')})
def dump(self) -> ArtifactCollection:
content = str(self.model.b).encode('utf-8')
yield Blobs({self.model_filename: InMemoryBlob(content)})
def eval_model(float_data):
return Model(
'model',
EvalModelWrapper(PickleModelIO()).bind_model('None',
input_data=float_data),
Blobs({}))
def artifact():
return Blobs({'kek': InMemoryBlob(b'kek')})
def dump(self) -> FilesContextManager:
yield Blobs({})
def write(self,
dataset: Dataset) -> Tuple[DatasetReader, ArtifactCollection]:
return TestDatasetReader(), Blobs(
{'data': InMemoryBlob(dataset.data.encode('utf8'))})
def dump(self) -> FilesContextManager:
model: lgb.Booster = self.model
with tempfile.TemporaryDirectory(prefix='ebonite_lightgbm_dump') as f:
path = os.path.join(f, self.model_path)
model.save_model(path)
yield Blobs({self.model_path: LocalFileBlob(path)})