def _append_repack_model_step()

in src/sagemaker/workflow/model_step.py [0:0]

• 62 lines of code
• 15 McCabe index (conditional complexity)

    def _append_repack_model_step(self):
        """Create and append a `_RepackModelStep` for the runtime repack"""
        if isinstance(self._model, PipelineModel):
            model_list = self._model.models
        elif isinstance(self._model, Model):
            model_list = [self._model]
        else:
            logger.warning("No models to repack")
            return

        self._pop_out_non_configurable_repack_model_step_args()

        security_group_ids, subnets = self._resolve_repack_model_step_vpc_configs()

        for i, model in enumerate(model_list):
            runtime_repack_flg = (
                self._need_runtime_repack and id(model) in self._need_runtime_repack
            )
            if runtime_repack_flg:
                name_base = model.name or i
                repack_model_step = _RepackModelStep(
                    name="{}-{}-{}".format(self.name, _REPACK_MODEL_NAME_BASE, name_base),
                    sagemaker_session=(
                        self._repack_model_step_settings.pop("sagemaker_session", None)
                        or self._model.sagemaker_session
                        or model.sagemaker_session
                    ),
                    role=(
                        self._repack_model_step_settings.pop("role", None)
                        or self._model.role
                        or model.role
                    ),
                    model_data=model.model_data,
                    entry_point=model.entry_point,
                    source_dir=model.source_dir,
                    dependencies=model.dependencies,
                    subnets=subnets,
                    security_group_ids=security_group_ids,
                    description=(
                        "Used to repack a model with customer scripts for a "
                        "register/create model step"
                    ),
                    depends_on=self.depends_on,
                    retry_policies=self._repack_model_retry_policies,
                    output_path=(
                        self._repack_model_step_settings.pop("output_path", None)
                        or self._runtime_repack_output_prefix
                    ),
                    output_kms_key=(
                        self._repack_model_step_settings.pop("output_kms_key", None)
                        or model.model_kms_key
                    ),
                    **self._repack_model_step_settings,
                )
                self.steps.append(repack_model_step)

                repacked_model_data = repack_model_step.properties.ModelArtifacts.S3ModelArtifacts
                if self._create_model_args:
                    if isinstance(self._model, PipelineModel):
                        container = self.step_args.create_model_request["Containers"][i]
                    else:
                        container = self.step_args.create_model_request["PrimaryContainer"]
                else:
                    container = self.step_args.create_model_package_request[
                        "InferenceSpecification"
                    ]["Containers"][i]
                container["ModelDataUrl"] = repacked_model_data