src/sagemaker/serve/model_server/djl_serving/prepare.py [28:85]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
logger = logging.getLogger(__name__)


def _extract_js_resource(js_model_dir: str, code_dir: Path, js_id: str):
    """Uncompress the jumpstart resource"""
    tmp_sourcedir = Path(js_model_dir).joinpath(f"infer-prepack-{js_id}.tar.gz")
    with tarfile.open(str(tmp_sourcedir)) as resources:
        custom_extractall_tarfile(resources, code_dir)


def _copy_jumpstart_artifacts(model_data: str, js_id: str, code_dir: Path) -> tuple:
    """Copy the associated JumpStart Resource into the code directory"""
    logger.info("Downloading JumpStart artifacts from S3...")

    s3_downloader = S3Downloader()
    if isinstance(model_data, str):
        if model_data.endswith(".tar.gz"):
            logger.info("Uncompressing JumpStart artifacts for faster loading...")
            with _tmpdir(directory=str(code_dir)) as js_model_dir:
                s3_downloader.download(model_data, js_model_dir)
                _extract_js_resource(js_model_dir, code_dir, js_id)
        else:
            logger.info("Copying uncompressed JumpStart artifacts...")
            s3_downloader.download(model_data, code_dir)
    elif (
        isinstance(model_data, dict)
        and model_data.get("S3DataSource")
        and model_data.get("S3DataSource").get("S3Uri")
    ):
        logger.info("Copying uncompressed JumpStart artifacts...")
        s3_downloader.download(model_data.get("S3DataSource").get("S3Uri"), code_dir)
    else:
        raise ValueError("JumpStart model data compression format is unsupported: %s", model_data)

    config_json_file = code_dir.joinpath("config.json")
    hf_model_config = None
    if config_json_file.is_file():
        with open(str(config_json_file)) as config_json:
            hf_model_config = json.load(config_json)

    return (hf_model_config, True)


def _create_dir_structure(model_path: str) -> tuple:
    """Placeholder Docstring"""
    model_path = Path(model_path)
    if not model_path.exists():
        model_path.mkdir(parents=True)
    elif not model_path.is_dir():
        raise ValueError("model_dir is not a valid directory")

    code_dir = model_path.joinpath("code")
    code_dir.mkdir(exist_ok=True, parents=True)

    _check_disk_space(model_path)
    _check_docker_disk_usage()

    return (model_path, code_dir)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -



src/sagemaker/serve/model_server/tgi/prepare.py [27:84]:
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
logger = logging.getLogger(__name__)


def _extract_js_resource(js_model_dir: str, code_dir: Path, js_id: str):
    """Uncompress the jumpstart resource"""
    tmp_sourcedir = Path(js_model_dir).joinpath(f"infer-prepack-{js_id}.tar.gz")
    with tarfile.open(str(tmp_sourcedir)) as resources:
        custom_extractall_tarfile(resources, code_dir)


def _copy_jumpstart_artifacts(model_data: str, js_id: str, code_dir: Path) -> tuple:
    """Copy the associated JumpStart Resource into the code directory"""
    logger.info("Downloading JumpStart artifacts from S3...")

    s3_downloader = S3Downloader()
    if isinstance(model_data, str):
        if model_data.endswith(".tar.gz"):
            logger.info("Uncompressing JumpStart artifacts for faster loading...")
            with _tmpdir(directory=str(code_dir)) as js_model_dir:
                s3_downloader.download(model_data, js_model_dir)
                _extract_js_resource(js_model_dir, code_dir, js_id)
        else:
            logger.info("Copying uncompressed JumpStart artifacts...")
            s3_downloader.download(model_data, code_dir)
    elif (
        isinstance(model_data, dict)
        and model_data.get("S3DataSource")
        and model_data.get("S3DataSource").get("S3Uri")
    ):
        logger.info("Copying uncompressed JumpStart artifacts...")
        s3_downloader.download(model_data.get("S3DataSource").get("S3Uri"), code_dir)
    else:
        raise ValueError("JumpStart model data compression format is unsupported: %s", model_data)

    config_json_file = code_dir.joinpath("config.json")
    hf_model_config = None
    if config_json_file.is_file():
        with open(str(config_json_file)) as config_json:
            hf_model_config = json.load(config_json)

    return (hf_model_config, True)


def _create_dir_structure(model_path: str) -> tuple:
    """Create the expected model directory structure for the TGI server"""
    model_path = Path(model_path)
    if not model_path.exists():
        model_path.mkdir(parents=True)
    elif not model_path.is_dir():
        raise ValueError("model_dir is not a valid directory")

    code_dir = model_path.joinpath("code")
    code_dir.mkdir(exist_ok=True, parents=True)

    _check_disk_space(model_path)
    _check_docker_disk_usage()

    return (model_path, code_dir)
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -