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)