def test_resolver_context(vineyard_client):
value = (1, 2, 3, 4)
o = vineyard_client.put(value)
result = vineyard_client.get(o)
assert result == value
assert result != 'faked tuple'
with resolver_context() as ctx:
ctx.register('vineyard::Tuple', fake_tuple_resolver)
result = vineyard_client.get(o)
assert result != value
assert result == 'faked tuple'
result = vineyard_client.get(o)
assert result == value
assert result != 'faked tuple'
with resolver_context({'vineyard::Tuple': fake_tuple_resolver}) as ctx:
result = vineyard_client.get(o)
assert result != value
assert result == 'faked tuple'
result = vineyard_client.get(o)
assert result == value
assert result != 'faked tuple'
def tf_dataframe_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
df = resolver(obj, **kw)
labels = df.pop(kw.get('label', 'label'))
if 'data' in kw:
return tf.data.Dataset.from_tensor_slices((np.stack(df[kw['data']], axis=0), labels))
return tf.data.Dataset.from_tensor_slices((dict(df), labels))
def mxnet_dataframe_resolver(obj, resolver, **kw):
with resolver_context(base=default_resolver_context) as resolver:
df = resolver(obj, **kw)
if 'label' in kw:
target = df[kw['label']].values.astype(np.float32)
data = df.drop(kw['label'], axis=1).values.astype(np.float32)
return mx.gluon.data.ArrayDataset((data, target))
def parallel_stream_resolver(obj):
"""Return a list of *local* partial streams."""
meta = obj.meta
partition_size = int(meta["size_"])
logger.debug('parallel stream: partitions = %d', partition_size)
with resolver_context() as ctx:
return [ctx(meta.get_member("stream_%d" % i)) for i in range(partition_size)]
def torch_dataframe_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
df = resolver(obj, **kw)
if 'label' in kw:
target = torch.tensor(df[kw['label']].values.astype(np.float32))
ds = torch.tensor(df.drop(kw['label'], axis=1).values.astype(np.float32))
return torch.utils.data.TensorDataset(ds, target)
def mxnet_table_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
table = resolver(obj, **kw)
table = table.to_pandas()
if 'label' in kw:
target = table[kw['label']].values.astype(np.float32)
data = table.drop(kw['label'], axis=1).values.astype(np.float32)
return mx.gluon.data.ArrayDataset((data, target))
def mxnet_record_batch_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
records = resolver(obj, **kw)
records = records.to_pandas()
if 'label' in kw:
target = records[kw['label']].values.astype(np.float32)
data = records.drop(kw['label'], axis=1).values.astype(np.float32)
return mx.gluon.data.ArrayDataset((data, target))
def torch_table_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
table = resolver(obj, **kw)
table = table.to_pandas()
if 'label' in kw:
target = torch.tensor(table[kw['label']].values)
ds = torch.tensor(table.drop(kw['label'], axis=1).values)
return torch.utils.data.TensorDataset(ds, target)
def torch_record_batch_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
records = resolver(obj, **kw)
records = records.to_pandas()
if 'label' in kw:
target = torch.tensor(records[kw['label']].values)
ds = torch.tensor(records.drop(kw['label'], axis=1).values)
return torch.utils.data.TensorDataset(ds, target)
def xgb_recordBatch_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
rb = resolver(obj, **kw)
# FIXME to_pandas is not zero_copy guaranteed
df = rb.to_pandas()
if 'label' in kw:
label = df.pop(kw['label'])
return xgb.DMatrix(df, label, feature_names=df.columns)
return xgb.DMatrix(df, feature_names=df.columns)
def xgb_dataframe_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
df = resolver(obj, **kw)
if 'label' in kw:
label = df.pop(kw['label'])
# data column can only be specified if label column is specified
if 'data' in kw:
df = np.stack(df[kw['data']].values)
return xgb.DMatrix(df, label)
return xgb.DMatrix(df, feature_names=df.columns)
def vineyard_for_pytorch():
with builder_context() as builder:
with resolver_context() as resolver:
register_torch_types(builder, resolver)
yield builder, resolver
def tf_record_batch_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
records = resolver(obj, **kw)
records = records.to_pandas()
labels = records.pop('label')
return tf.data.Dataset.from_tensor_slices((dict(records), labels))
def vineyard_for_dask():
with builder_context() as builder:
with resolver_context() as resolver:
register_dask_types(builder, resolver)
yield builder, resolver
def xgb_tensor_resolver(obj, **kw):
with resolver_context(base=default_resolver_context) as resolver:
array = resolver(obj, **kw)
return xgb.DMatrix(array)
def vineyard_for_tensorflow():
with builder_context() as builder:
with resolver_context() as resolver:
register_tf_types(builder, resolver)
yield builder, resolver
