Support removing task input models using Task.remove_input_models()

This commit is contained in:
allegroai 2023-05-21 09:41:00 +03:00
parent fa58c69fc3
commit e80d1f1ff4

View File

@ -56,6 +56,9 @@ from .repo import ScriptInfo, pip_freeze
from .hyperparams import HyperParams
from ...config import config, PROC_MASTER_ID_ENV_VAR, SUPPRESS_UPDATE_MESSAGE_ENV_VAR, DOCKER_BASH_SETUP_ENV_VAR
from ...utilities.process.mp import SingletonLock
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from ...model import BaseModel
class Task(IdObjectBase, AccessMixin, SetupUploadMixin):
@ -1374,6 +1377,22 @@ class Task(IdObjectBase, AccessMixin, SetupUploadMixin):
execution.model_labels = enumeration
self._edit(execution=execution)
def remove_input_models(self, models_to_remove):
# type: (Sequence[Union[str, BaseModel]]) -> ()
"""
Remove input models from the current task. Note that the models themselves are not deleted,
but the tasks' reference to the models is removed.
To delete the models themselves, see `Models.remove`
:param models_to_remove: The models to remove from the task. Can be a list of ids,
or of `BaseModel` (including its subclasses: `Model` and `InputModel`)
"""
ids_to_remove = [model if isinstance(model, str) else model.id for model in models_to_remove]
with self._edit_lock:
self.reload()
self.data.models.input = [model for model in self.data.models.input if model.model not in ids_to_remove]
self._edit(models=self.data.models)
def _set_default_docker_image(self):
# type: () -> ()
if not DOCKER_IMAGE_ENV_VAR.exists() and not DOCKER_BASH_SETUP_ENV_VAR.exists():