Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_tar_and_upload_dir_with_subdirectory(sagemaker_session, tmpdir):
file_tree(tmpdir, ["src-dir/sub/train.py"])
source_dir = os.path.join(str(tmpdir), "src-dir")
with patch("shutil.rmtree"):
result = fw_utils.tar_and_upload_dir(
sagemaker_session, "bucket", "prefix", "train.py", source_dir
)
assert result == fw_utils.UploadedCode(
s3_prefix="s3://bucket/prefix/sourcedir.tar.gz", script_name="train.py"
)
assert {"/sub/train.py"} == list_source_dir_files(sagemaker_session, tmpdir)
def test_tar_and_upload_dir_no_directory_bare_filename(sagemaker_session, tmpdir):
source_dir = file_tree(tmpdir, ["train.py"])
entrypoint = "train.py"
with patch("shutil.rmtree"):
with cd(source_dir):
result = fw_utils.tar_and_upload_dir(
sagemaker_session, "bucket", "prefix", entrypoint, None
)
assert result == fw_utils.UploadedCode(
s3_prefix="s3://bucket/prefix/sourcedir.tar.gz", script_name="train.py"
)
assert {"/train.py"} == list_source_dir_files(sagemaker_session, tmpdir)
if repack:
bucket = self.bucket or self.sagemaker_session.default_bucket()
repacked_model_data = "s3://" + os.path.join(bucket, key_prefix, "model.tar.gz")
utils.repack_model(
inference_script=self.entry_point,
source_directory=self.source_dir,
dependencies=self.dependencies,
model_uri=self.model_data,
repacked_model_uri=repacked_model_data,
sagemaker_session=self.sagemaker_session,
kms_key=self.model_kms_key,
)
self.repacked_model_data = repacked_model_data
self.uploaded_code = UploadedCode(
s3_prefix=self.repacked_model_data, script_name=os.path.basename(self.entry_point)
)
job_name (str): The new job name included in the submit S3 URI
Returns:
str: The updated S3 URI of framework source directory
"""
if estimator.uploaded_code is None:
return
pattern = r"(?<=/)[^/]+?(?=/source/sourcedir.tar.gz)"
# update the S3 URI with the latest training job.
# s3://path/old_job/source/sourcedir.tar.gz will become s3://path/new_job/source/sourcedir.tar.gz
submit_uri = estimator.uploaded_code.s3_prefix
submit_uri = re.sub(pattern, job_name, submit_uri)
script_name = estimator.uploaded_code.script_name
estimator.uploaded_code = fw_utils.UploadedCode(submit_uri, script_name)
model.py_version,
)
base_name = utils.base_name_from_image(deploy_image)
model.name = model.name or utils.name_from_base(base_name)
bucket = model.bucket or model.sagemaker_session._default_bucket
script = os.path.basename(model.entry_point)
key = "{}/source/sourcedir.tar.gz".format(model.name)
if model.source_dir and model.source_dir.lower().startswith("s3://"):
code_dir = model.source_dir
model.uploaded_code = fw_utils.UploadedCode(s3_prefix=code_dir, script_name=script)
else:
code_dir = "s3://{}/{}".format(bucket, key)
model.uploaded_code = fw_utils.UploadedCode(s3_prefix=code_dir, script_name=script)
s3_operations["S3Upload"] = [
{"Path": model.source_dir or script, "Bucket": bucket, "Key": key, "Tar": True}
]
deploy_env = dict(model.env)
deploy_env.update(model._framework_env_vars())
try:
if model.model_server_workers:
deploy_env[sagemaker.model.MODEL_SERVER_WORKERS_PARAM_NAME.upper()] = str(
model.model_server_workers
)
except AttributeError:
# This applies to a FrameworkModel which is not SageMaker Deep Learning Framework Model
pass
bucket, key = fw_utils.parse_s3_url(estimator.code_location)
key = os.path.join(key, estimator._current_job_name, "source", "sourcedir.tar.gz")
elif estimator.uploaded_code is not None:
bucket, key = fw_utils.parse_s3_url(estimator.uploaded_code.s3_prefix)
else:
bucket = estimator.sagemaker_session._default_bucket
key = os.path.join(estimator._current_job_name, "source", "sourcedir.tar.gz")
script = os.path.basename(estimator.entry_point)
if estimator.source_dir and estimator.source_dir.lower().startswith("s3://"):
code_dir = estimator.source_dir
estimator.uploaded_code = fw_utils.UploadedCode(s3_prefix=code_dir, script_name=script)
else:
code_dir = "s3://{}/{}".format(bucket, key)
estimator.uploaded_code = fw_utils.UploadedCode(s3_prefix=code_dir, script_name=script)
s3_operations["S3Upload"] = [
{
"Path": estimator.source_dir or estimator.entry_point,
"Bucket": bucket,
"Key": key,
"Tar": True,
}
]
estimator._hyperparameters[sagemaker.model.DIR_PARAM_NAME] = code_dir
estimator._hyperparameters[sagemaker.model.SCRIPT_PARAM_NAME] = script
estimator._hyperparameters[
sagemaker.model.CLOUDWATCH_METRICS_PARAM_NAME
] = estimator.enable_cloudwatch_metrics
estimator._hyperparameters[
sagemaker.model.CONTAINER_LOG_LEVEL_PARAM_NAME
] = estimator.container_log_level