def test_save_dict_as_pkl_s3(self, s3_fs, sample_dict):
test_path = f"s3://{test_bucket_name}/tests/dict/dict.pkl"
save_path = f"s3://{test_bucket_name}/copy/dict.pkl"
io.save_object(sample_dict, save_path, file_type="pickle", fs=s3_fs)
test_obj = pickle.loads(s3_fs.cat(test_path))
save_obj = pickle.loads(s3_fs.cat(save_path))
assert test_obj == save_obj
def test_save_dict_as_json_s3(self, s3_fs, sample_dict):
test_path = f"s3://{test_bucket_name}/tests/dict/dict.json"
save_path = f"s3://{test_bucket_name}/copy/dict.json"
io.save_object(sample_dict, save_path, file_type="json", fs=s3_fs)
test_obj = s3_fs.cat(test_path)
save_obj = s3_fs.cat(save_path)
assert test_obj == save_obj
def test_save_dict_as_json_local(self, sample_local_dir, sample_dict):
test_path = os.path.join(sample_local_dir, "io_tests/dict/dict.json")
save_path = os.path.join(sample_local_dir, "io_tests/copy/dict.json")
io.save_object(sample_dict, save_path, file_type="json")
with open(test_path, "r") as f:
test_obj = f.read()
with open(save_path, "r") as f:
save_obj = f.read()
assert test_obj == save_obj
def test_save_dict_as_pkl_local(self, sample_local_dir, sample_dict):
test_path = os.path.join(sample_local_dir, "io_tests/dict/dict.pkl")
save_path = os.path.join(sample_local_dir, "io_tests/copy/dict.pkl")
io.save_object(sample_dict, save_path, file_type="pickle")
with open(test_path, "rb") as f:
test_obj = pickle.load(f)
with open(save_path, "rb") as f:
save_obj = pickle.load(f)
assert test_obj == save_obj
def save_artifact(obj, subpath: str, **kwargs) -> None:
""" Saves obj to the specified subpath
Parameters
-----------
obj : object
Python object in memory
subpath : str
The location to save the object to. The subpath will be placed within
the currently active mlflow run_uuid within the specified artifact
store (this can be specified when creating the mlflow run)
kwargs : Dict
All other arguments are passed to io.save_object
Returns
--------
None
"""
path = os.path.join(mlflow_.get_artifact_uri(), subpath)
logger.info(f"Saving artifact to {path}")
io.save_object(obj, path, **kwargs)
def test_save_invalid_file_type(self):
obj = [1, 2, 3]
path = f"s3://{test_bucket_name}/copy/foo"
with pytest.raises(ValueError):
io.save_object(obj, path, file_type="foobar", fs=s3_fs)