def _download_and_extract_model_tar_gz()

in src/vw-serving/src/vw_serving/model_manager.py [0:0]

• 19 lines of code
• 3 McCabe index (conditional complexity)

    def _download_and_extract_model_tar_gz(self, model_id):
        """
        This function first gets the s3 location from dynamo db,
        downloads the model, extracts it and then
        returns a tuple (str, str) of metadata string and model weights URL on disk
        """
        deployable_model_id_record = self.model_ddb_wrapper.get_model_record(experiment_id=self.experiment_id,
                                                                             model_id=model_id)
        s3_uri = deployable_model_id_record.get("s3_model_output_path", "")
        if s3_uri:
            try:
                tmp_dir = Path(f"/opt/ml/downloads/{gen_random_string()}")
                tmp_dir.mkdir(parents=True, exist_ok=True)
                tmp_model_tar_gz = os.path.join(tmp_dir.as_posix(), "model.tar.gz")
                bucket, key = parse_s3_url(s3_uri)
                self.s3_resource.Bucket(bucket).download_file(key, tmp_model_tar_gz)
                shutil.unpack_archive(filename=tmp_model_tar_gz, extract_dir=tmp_dir.as_posix())
                return self.get_model(tmp_dir.as_posix())
            except Exception as e:
                logger.exception(f"Could not parse or download {model_id} from {s3_uri} due to {e}")
                return None
        else:
            logger.exception(f"Could not s3 location of {model_id}")
            return None