Skip to content

Commit

Permalink
fix: upload models for bare push bento
Browse files Browse the repository at this point in the history
Signed-off-by: Frost Ming <me@frostming.com>
  • Loading branch information
frostming committed Aug 30, 2024
1 parent 7407358 commit be2f7a2
Show file tree
Hide file tree
Showing 2 changed files with 41 additions and 42 deletions.
37 changes: 19 additions & 18 deletions src/bentoml/_internal/bento/bento.py
Original file line number Diff line number Diff line change
Expand Up @@ -282,26 +282,27 @@ def create(
temp_dir=BentoMLContainer.tmp_bento_store_dir.get(),
)
models: list[BentoModelInfo] = []
if not bare:
ctx_fs = fs.open_fs(encode_path_for_uri(build_ctx))

def append_model(model: BentoModelInfo) -> None:
if model not in models:
models.append(model)

if build_config.models:
for model_spec in build_config.models:
model = BentoModel(model_spec.tag)
append_model(model.to_info(model_spec.alias))
elif is_legacy:
# XXX: legacy way to get models from service
# Add all models required by the service
for model in svc.models:
def append_model(model: BentoModelInfo) -> None:
if model not in models:
models.append(model)

if build_config.models:
for model_spec in build_config.models:
model = BentoModel(model_spec.tag)
append_model(model.to_info(model_spec.alias))
elif is_legacy:
# XXX: legacy way to get models from service
# Add all models required by the service
for model in svc.models:
append_model(BentoModel(model.tag).to_info())
# Add all models required by service runners
for runner in svc.runners:
for model in runner.models:
append_model(BentoModel(model.tag).to_info())
# Add all models required by service runners
for runner in svc.runners:
for model in runner.models:
append_model(BentoModel(model.tag).to_info())

if not bare:
ctx_fs = fs.open_fs(encode_path_for_uri(build_ctx))

# create ignore specs
specs = BentoPathSpec(build_config.include, build_config.exclude)
Expand Down
46 changes: 22 additions & 24 deletions src/bentoml/_internal/cloud/bentocloud.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,31 +94,29 @@ def _do_push_bento(
if version is None:
raise BentoMLException(f"Bento {bento.tag} version cannot be None")
info = bento.info
if not bare:
models_to_push: list[Model[t.Any]] = []
for model in info.all_models:
if model.registry == "huggingface":
models_to_push.append(HuggingFaceModel(model.tag, model.endpoint))
else:
model = BentoModel(model.tag)
if model.stored is not None:
models_to_push.append(model)
with ThreadPoolExecutor(
max_workers=max(len(models_to_push), 1)
) as executor:

def push_model(model: BentoModel) -> None:
model_upload_task_id = self.spinner.transmission_progress.add_task(
f'Pushing model "{model.tag}"', start=False, visible=False
)
self._do_push_model(
model,
model_upload_task_id,
force=force,
threads=threads,
)
# Upload models
models_to_push: list[Model[t.Any]] = []
for model in info.all_models:
if model.registry == "huggingface":
models_to_push.append(HuggingFaceModel(model.tag, model.endpoint))
else:
model = BentoModel(model.tag)
if model.stored is not None:
models_to_push.append(model)
with ThreadPoolExecutor(max_workers=max(len(models_to_push), 1)) as executor:

def push_model(model: BentoModel) -> None:
model_upload_task_id = self.spinner.transmission_progress.add_task(
f'Pushing model "{model.tag}"', start=False, visible=False
)
self._do_push_model(
model,
model_upload_task_id,
force=force,
threads=threads,
)

executor.map(push_model, models_to_push)
executor.map(push_model, models_to_push)
with self.spinner.spin(text=f'Fetching Bento repository "{name}"'):
bento_repository = rest_client.v1.get_bento_repository(
bento_repository_name=name
Expand Down

0 comments on commit be2f7a2

Please sign in to comment.