def enable_sm_oneclick_deploy()

in sagemaker/sagemaker-inference/pytorch/code_mme/train.py [0:0]

• 18 lines of code
• 8 McCabe index (conditional complexity)

def enable_sm_oneclick_deploy(model_dir):
    """Copy current running source code folder to model_dir, to enable Estimator.deploy()
    PyTorch framework containers will load custom inference code if:
    - The code exists in a top-level code/ folder in the model.tar.gz
    - The entry point argument matches an existing file
    ...So to make one-click estimator.deploy() work (without creating a PyTorchModel first), we need
    to:
    - Copy the current working directory to model_dir/code
    - `from inference import *` because "train.py" will still be the entry point (same as the training job)
    """
    code_path = os.path.join(model_dir, "code")
    logger.info(f"Copying working folder to {code_path}")
    for currpath, dirs, files in os.walk("."):
        for file in files:
            # Skip any filenames starting with dot:
            if file.startswith("."):
                continue
            filepath = os.path.join(currpath, file)
            # Skip any pycache or dot folders:
            if ((os.path.sep + ".") in filepath) or ("__pycache__" in filepath):
                continue
            relpath = filepath[len(".") :]
            if relpath.startswith(os.path.sep):
                relpath = relpath[1:]
            outpath = os.path.join(code_path, relpath)
            logger.info(f"Copying {filepath} to {outpath}")
            os.makedirs(outpath.rpartition(os.path.sep)[0], exist_ok=True)
            shutil.copy2(filepath, outpath)
    return code_path