def _upload_imzmls_from_prefix_if_needed(src_path,
storage,
sm_storage,
s3_client=None):
if src_path.startswith('cos://'):
bucket, prefix = src_path[len('cos://'):].split('/', maxsplit=1)
keys = [
f'cos://{bucket}/{key}'
for key in storage.list_keys(bucket, prefix)
]
elif src_path.startswith('s3a://'):
bucket, prefix = split_s3_path(src_path)
response = s3_client.list_objects_v2(Bucket=bucket, Prefix=prefix)
if 'Contents' in response:
keys = [
f"s3a://{bucket}/{item['Key']}"
for item in response['Contents']
]
else:
keys = []
else:
keys = [str(p) for p in Path(src_path).iterdir()]
imzml_keys = [key for key in keys if key.lower().endswith('.imzml')]
ibd_keys = [key for key in keys if key.lower().endswith('.ibd')]
assert len(imzml_keys) == 1, imzml_keys
assert len(ibd_keys) == 1, ibd_keys
imzml_cobj = _upload_if_needed(imzml_keys[0], storage, sm_storage, 'imzml',
s3_client)
ibd_cobj = _upload_if_needed(ibd_keys[0], storage, sm_storage, 'imzml',
s3_client)
return imzml_cobj, ibd_cobj
def parse_input_path_for_lithops(sm_config, input_path):
if input_path.startswith('s3://') or input_path.startswith('s3a://'):
backend = 'aws_s3'
bucket, prefix = split_s3_path(input_path)
else:
backend = 'ibm_cos'
bucket, prefix = split_cos_path(input_path)
storage = Storage(sm_config['lithops'], backend)
if backend == 'aws_s3' and sm_config['lithops']['aws_s3'][
'endpoint'].startswith('http://'):
# WORKAROUND for local Minio access
# Lithops forces the url to HTTPS, so overwrite the S3 client with a fixed client
# https://github.com/lithops-cloud/lithops/issues/708
storage.storage_handler.s3_client = get_s3_client()
keys_in_path = storage.list_keys(bucket, prefix)
imzml_keys = [
key for key in keys_in_path if key.lower().endswith('.imzml')
]
ibd_keys = [key for key in keys_in_path if key.lower().endswith('.ibd')]
debug_info = f'Path {input_path} had keys: {keys_in_path}'
assert len(
imzml_keys) == 1, f'Couldn\'t determine imzML file. {debug_info}'
assert len(ibd_keys) == 1, f'Couldn\'t determine ibd file. {debug_info}'
imzml_cobject = CloudObject(storage.backend, bucket, imzml_keys[0])
ibd_cobject = CloudObject(storage.backend, bucket, ibd_keys[0])
return storage, imzml_cobject, ibd_cobject
def read_moldb_file(file_path):
try:
if re.findall(r'^s3a?://', file_path):
bucket_name, key = split_s3_path(file_path)
sm_config = SMConfig.get_conf()
buffer = get_s3_bucket(bucket_name,
sm_config).Object(key).get()['Body']
else:
buffer = Path(file_path).open()
moldb_df = pd.read_csv(buffer, sep='\t', dtype=object, na_filter=False)
except ValueError as e:
raise MalformedCSV(f'Malformed CSV: {e}') from e
if moldb_df.empty:
raise MalformedCSV('No data rows found')
required_columns = {'id', 'name', 'formula'}
if not required_columns.issubset(set(moldb_df.columns)):
raise MalformedCSV(
f'Missing columns. Provided: {moldb_df.columns.to_list()} Required: {required_columns}'
)
parsing_errors = _validate_moldb_df(moldb_df)
if parsing_errors:
raise BadData('Failed to parse some rows', *parsing_errors)
moldb_df.rename({
'id': 'mol_id',
'name': 'mol_name'
},
axis='columns',
inplace=True)
return moldb_df
def load_scoring_model(name: Optional[str]) -> ScoringModel:
# Import DB locally so that Lithops doesn't try to pickle it & fail due to psycopg2
# pylint: disable=import-outside-toplevel # circular import
from sm.engine.db import DB
if name is None:
return MsmScoringModel()
row = DB().select_one(
"SELECT type, params FROM scoring_model WHERE name = %s", (name, ))
assert row, f'Scoring model {name} not found'
type_, params = row
if type_ == 'catboost':
bucket, key = split_s3_path(params['s3_path'])
with TemporaryDirectory() as tmpdir:
model_file = Path(tmpdir) / 'model.cbm'
with model_file.open('wb') as f:
f.write(get_s3_client().get_object(Bucket=bucket,
Key=key)['Body'].read())
model = CatBoost()
model.load_model(str(model_file), 'cbm')
return CatBoostScoringModel(name, model, params)
else:
raise ValueError(f'Unsupported scoring model type: {type_}')
def exists(self, path):
try:
self.s3.Object(*split_s3_path(path)).load()
except ClientError as e:
if e.response['Error']['Code'] == "404":
return False
else:
raise e
else:
logger.info('Path s3://%s/%s already exists', self.bucket, path)
return True
def _saved(self):
"""Check if ion centroids saved to parquet"""
if self._centroids_stored_on_s3:
bucket, key = split_s3_path(self._ion_centroids_path)
try:
for fn in self._parquet_file_names:
self._s3.head_object(Bucket=bucket, Key=f'{key}/{fn}')
except ClientError:
return False
else:
return True
else:
return all(
(Path(self._ion_centroids_path) / fn).exists() for fn in self._parquet_file_names
)
def _copy_input_data(self, ds):
logger.info('Copying input data')
self._ds_data_path = Path(self._sm_config['fs']['spark_data_path']) / ds.id
if ds.input_path.startswith('s3a://'):
self._ds_data_path.mkdir(parents=True, exist_ok=True)
bucket_name, key = split_s3_path(ds.input_path)
bucket = storage.get_s3_bucket(bucket_name, self._sm_config)
for obj_sum in bucket.objects.filter(Prefix=key):
local_file = str(self._ds_data_path / Path(obj_sum.key).name)
logger.debug(f'Downloading s3a://{bucket_name}/{obj_sum.key} -> {local_file}')
obj_sum.Object().download_file(local_file)
else:
rmtree(self._ds_data_path, ignore_errors=True)
copytree(src=ds.input_path, dst=self._ds_data_path)
def _upload_if_needed(src_path,
storage,
sm_storage,
storage_type,
s3_client=None,
use_db_mutex=True):
"""
Uploads the object from `src_path` if it doesn't already exist in its translated COS path.
Returns a CloudObject for the COS object
"""
bucket, key = _choose_cos_location(src_path, sm_storage, storage_type)
with ExitStack() as stack:
if use_db_mutex:
# Lock during upload to prevent parallel jobs upload the same file simultaneously
stack.enter_context(DBMutex().lock(bucket + key, timeout=1200))
try:
storage.head_object(bucket, key)
logger.debug(f'{src_path} already uploaded')
return CloudObject(storage.backend, bucket, key)
except StorageNoSuchKeyError:
logger.info(f'Uploading {src_path}...')
if src_path.startswith('s3a://'):
assert s3_client, 'S3 client must be supplied to support s3a:// paths'
src_bucket, src_key = split_s3_path(src_path)
obj = s3_client.get_object(Bucket=src_bucket, Key=src_key)
if hasattr(storage.get_client(), 'upload_fileobj'):
# Try streaming upload to IBM COS
transfer_config = TransferConfig(multipart_chunksize=20 *
MB,
max_concurrency=20,
io_chunksize=1 * MB)
storage.get_client().upload_fileobj(Fileobj=obj['Body'],
Bucket=bucket,
Key=key,
Config=transfer_config)
cobject = CloudObject(storage.backend, bucket, key)
else:
# Fall back to buffering the entire object in memory for other backends
cobject = storage.put_cloudobject(obj['Body'].read(),
bucket, key)
else:
cobject = storage.put_cloudobject(open(src_path, 'rb'), bucket,
key)
logger.info(f'Uploading {src_path}...Done')
return cobject
def exists(self):
""" Check if ion centroids saved to parquet
"""
if self._ion_centroids_path.startswith('s3a://'):
cred_dict = dict(aws_access_key_id=self._sm_config['aws']['aws_access_key_id'],
aws_secret_access_key=self._sm_config['aws']['aws_secret_access_key'])
bucket, key = split_s3_path(self._ion_centroids_path)
s3 = boto3.client('s3', **cred_dict)
try:
s3.head_object(Bucket=bucket, Key=key + '/ions/_SUCCESS')
except ClientError:
return False
else:
return True
else:
return Path(self._ion_centroids_path + '/ions/_SUCCESS').exists()
def _choose_cos_location(src_path, sm_storage, storage_type):
"""Maps the provided COS/S3/local filesystem path to an appropriate bucket & key in COS"""
bucket, prefix = sm_storage[storage_type]
src_path = str(src_path)
if src_path.startswith('cos://'):
# Already in COS - no need to translate path
return split_cos_path(src_path)
if src_path.startswith('s3a://'):
# Ignore the bucket and take the key
_, suffix = split_s3_path(src_path)
else:
# Ignore the directory and take the filename
suffix = Path(src_path).name
key = f'{prefix}/{suffix}' if prefix else suffix
return bucket, key
def exists(self):
""" Check if ion centroids saved to parquet
"""
if self._ion_centroids_path.startswith('s3a://'):
cred_dict = dict(
aws_access_key_id=self._sm_config['aws']['aws_access_key_id'],
aws_secret_access_key=self._sm_config['aws']
['aws_secret_access_key'])
bucket, key = split_s3_path(self._ion_centroids_path)
s3 = boto3.client('s3', **cred_dict)
try:
s3.head_object(Bucket=bucket, Key=key + '/ions/_SUCCESS')
except ClientError:
return False
else:
return True
else:
return Path(self._ion_centroids_path + '/ions/_SUCCESS').exists()
def copy_input_data(self, input_data_path):
""" Copy mass spec files from input path to a dataset work directory
Args
----
input_data_path : str
Path to input files
"""
logger.info('Copying data from %s to %s', input_data_path, self.local_dir.ds_path)
if input_data_path.startswith('s3a://'):
cmd_check('mkdir -p {}', self.local_dir.ds_path)
bucket_name, inp_path = split_s3_path(input_data_path)
bucket = self.s3.Bucket(bucket_name)
for obj in bucket.objects.filter(Prefix=inp_path):
if not obj.key.endswith('/'):
path = join(self.local_dir.ds_path, obj.key.split('/')[-1])
self.s3transfer.download_file(bucket_name, obj.key, path)
else:
self.local_dir.copy(input_data_path, self.local_dir.ds_path)
# MZ FILE SEARCH ON S3
#
import io
import boto3
from sm.engine import util as sm_engine_utils
from sm.browser import mz_search
s3 = boto3.resource("s3")
list(s3.buckets.all())
s3_path = (
"s3://sm-engine-dev/dataset-browser/"
"20200228_366x629_30um_Mouse_Obesity_16w_DAN_Neg_mode_190-2000mz_70K_Laser37_6/"
"segment_0000.bin"
)
bucket_name, key = sm_engine_utils.split_s3_path(s3_path)
mz_bin_object = s3.Object(bucket_name=bucket_name, key=key)
type(mz_bin_object), mz_bin_object.content_length
# %time _ = mz_bin_object.content_length
s3_file = mz_search.S3File(mz_bin_object)
# s3_file.seek(0)
# bytes = s3_file.read(12)
# # %time bytes = s3_file.read(12)
# mz_chunks_array = np.frombuffer(bytes, dtype="f").reshape(-1, 3)
mz = 231.0927
ppm = 3
mz_lo, mz_hi = mz - mz * ppm * 1e-6, mz + mz * ppm * 1e-6
import random
def _upload_to_s3(self):
bucket, key = split_s3_path(self._ion_centroids_path)
for fn in self._parquet_file_names:
self._s3.upload_file(
Filename=str(self._local_ion_centroids_path / fn), Bucket=bucket, Key=f'{key}/{fn}'
)
def copy(self, local, remote):
logger.info('Coping from {} to {}'.format(local, remote))
self.s3transfer.upload_file(local, *split_s3_path(remote))
def _download_from_s3(self):
bucket, key = split_s3_path(self._ion_centroids_path)
for fn in self._parquet_file_names:
self._s3.download_file(
Bucket=bucket, Key=f'{key}/{fn}', Filename=str(self._local_ion_centroids_path / fn)
)
def __init__(self, base_path, ds_id, s3, s3transfer):
self.s3 = s3
self.s3transfer = s3transfer
self.bucket, path = split_s3_path(base_path)
self.ds_path = join(path, ds_id)
def del_input_data(self, input_data_path):
if input_data_path.startswith('s3a://'):
bucket, path = split_s3_path(input_data_path)
delete_s3_path(bucket, path, self.s3)
else:
delete_local_path(input_data_path)