def test_repack_entry_point_only(tmp):
model_name = "xg-boost-model"
fake_model_path = os.path.join(tmp, model_name)
# create a fake model
open(fake_model_path, "w")
# create model.tar.gz
model_tar_path = "s3://my-bucket/model-%s.tar.gz" % time.time()
model_tar_name = model_tar_path.split("/")[-1]
model_tar_location = os.path.join(tmp, model_tar_name)
with tarfile.open(model_tar_location, mode="w:gz") as t:
t.add(fake_model_path, arcname=model_name)
# move model.tar.gz to /opt/ml/input/data/training
Path("/opt/ml/input/data/training").mkdir(parents=True, exist_ok=True)
shutil.move(model_tar_location,
os.path.join("/opt/ml/input/data/training", model_tar_name))
# create files that will be added to model.tar.gz
create_file_tree(
"/opt/ml/code",
[
"inference.py",
],
)
# repack
_repack_model.repack(inference_script="inference.py",
model_archive=model_tar_path)
# /opt/ml/model should now have the original model and the inference script
assert os.path.exists(os.path.join("/opt/ml/model", model_name))
assert os.path.exists(os.path.join("/opt/ml/model/code", "inference.py"))
def test_repack_with_source_dir_and_dependencies(tmp):
model_name = "xg-boost-model"
fake_model_path = os.path.join(tmp, model_name)
# create a fake model
open(fake_model_path, "w")
# create model.tar.gz
model_tar_name = "model-%s.tar.gz" % time.time()
model_tar_location = os.path.join(tmp, model_tar_name)
with tarfile.open(model_tar_location, mode="w:gz") as t:
t.add(fake_model_path, arcname=model_name)
# move model.tar.gz to /opt/ml/input/data/training
Path("/opt/ml/input/data/training").mkdir(parents=True, exist_ok=True)
shutil.move(model_tar_location,
os.path.join("/opt/ml/input/data/training", model_tar_name))
# create files that will be added to model.tar.gz
create_file_tree(
"/opt/ml/code",
[
"inference.py",
"dependencies/a",
"bb",
"dependencies/some/dir/b",
"sourcedir/foo.py",
"sourcedir/some/dir/a",
],
)
# repack
_repack_model.repack(
inference_script="inference.py",
model_archive=model_tar_name,
dependencies="dependencies/a bb dependencies/some/dir",
source_dir="sourcedir",
)
# /opt/ml/model should now have the original model and the inference script
assert os.path.exists(os.path.join("/opt/ml/model", model_name))
assert os.path.exists(os.path.join("/opt/ml/model/code", "inference.py"))
assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "a"))
assert os.path.exists(os.path.join("/opt/ml/model/code/lib", "bb"))
assert os.path.exists(os.path.join("/opt/ml/model/code/lib/dir", "b"))
assert os.path.exists(os.path.join("/opt/ml/model/code/", "foo.py"))
assert os.path.exists(os.path.join("/opt/ml/model/code/some/dir", "a"))