def testBearerTokenAccount(self):
self.odps.delete_table(tn('test_bearer_token_account_table'),
if_exists=True)
t = self.odps.create_table(tn('test_bearer_token_account_table'),
'col string',
lifecycle=1)
with t.open_writer() as writer:
records = [['val1'], ['val2'], ['val3']]
writer.write(records)
inst = self.odps.execute_sql('select count(*) from {0}'.format(
tn('test_bearer_token_account_table')),
async_=True)
inst.wait_for_success()
task_name = inst.get_task_names()[0]
logview_address = inst.get_logview_address()
token = logview_address[logview_address.find('token=') +
len('token='):]
bearer_token_account = BearerTokenAccount(token=token)
bearer_token_odps = ODPS(None,
None,
self.odps.project,
self.odps.endpoint,
account=bearer_token_account)
bearer_token_instance = bearer_token_odps.get_instance(inst.id)
self.assertEqual(inst.get_task_result(task_name),
bearer_token_instance.get_task_result(task_name))
self.assertEqual(inst.get_task_summary(task_name),
bearer_token_instance.get_task_summary(task_name))
with self.assertRaises(errors.NoPermission):
bearer_token_odps.create_table(
tn('test_bearer_token_account_table_test1'),
'col string',
lifecycle=1)
fake_token_account = BearerTokenAccount(token='fake-token')
bearer_token_odps = ODPS(None,
None,
self.odps.project,
self.odps.endpoint,
account=fake_token_account)
with self.assertRaises(errors.ODPSError):
bearer_token_odps.create_table(
tn('test_bearer_token_account_table_test2'),
'col string',
lifecycle=1)
def _execute_in_cupid(cls, ctx, op):
import os
import pandas as pd
from odps import ODPS
from odps.accounts import BearerTokenAccount
cupid_client = CupidServiceClient()
to_store_data = ctx[op.inputs[0].key]
bearer_token = cupid_client.get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME', None)
odps_params = op.odps_params.copy()
if project:
odps_params['project'] = project
endpoint = os.environ.get(
'ODPS_RUNTIME_ENDPOINT') or odps_params['endpoint']
o = ODPS(None,
None,
account=account,
project=odps_params['project'],
endpoint=endpoint)
odps_schema = o.get_table(op.table_name).schema
project_name, table_name = op.table_name.split('.')
writer_config = dict(_table_name=table_name,
_project_name=project_name,
_table_schema=odps_schema,
_partition_spec=op.partition_spec,
_block_id=op.block_id,
_handle=op.cupid_handle)
cupid_client.write_table_data(writer_config, to_store_data,
op.write_batch_size)
ctx[op.outputs[0].key] = pd.DataFrame()
def check_instance_idle(self):
last_active_time = self._last_activity_time
has_running = False
for ref in self._session_refs.values():
for info in ref.get_graph_infos().values():
if info.get('end_time') is None:
has_running = True
break
else:
last_active_time = max(info['end_time'], last_active_time)
if has_running:
break
if not has_running and last_active_time < time.time(
) - self._idle_timeout:
# timeout: we need to kill the instance
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid.runtime import context
logger.warning('Timeout met, killing the instance now.')
bearer_token = context().get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ['ODPS_PROJECT_NAME']
endpoint = os.environ['ODPS_RUNTIME_ENDPOINT']
o = ODPS(None,
None,
account=account,
project=project,
endpoint=endpoint)
o.stop_instance(os.environ['MARS_K8S_POD_NAMESPACE'])
else:
self.ref().check_instance_idle(_delay=10, _tell=True, _wait=False)
def _handle_terminate_instance(sock):
from cupid.runtime import context, RuntimeContext
from odps import ODPS
from odps.accounts import BearerTokenAccount
try:
cmd_len, = struct.unpack('<I', sock.recv(4))
# dict with key
cmd_body = pickle.loads(sock.recv(cmd_len))
instance_id = cmd_body['instance_id']
if not RuntimeContext.is_context_ready():
logger.warning('Cupid context not ready')
else:
bearer_token = context().get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ['ODPS_PROJECT_NAME']
endpoint = os.environ['ODPS_RUNTIME_ENDPOINT']
o = ODPS(None,
None,
account=account,
project=project,
endpoint=endpoint)
o.stop_instance(instance_id)
except:
logger.exception('Failed to put kv value')
_write_request_result(sock, False, exc_info=sys.exc_info())
def check_instance_idle(self):
from cupid.runtime import context
has_running, active_time_from_service = self._get_service_activity_info(
)
if active_time_from_service != self._last_active_time_from_service:
self._last_active_time = active_time_from_service
self._last_active_time_from_service = active_time_from_service
elif has_running:
self._last_active_time = time.time()
if self._last_active_time < time.time() - self._idle_timeout:
# timeout: we need to kill the instance
from odps import ODPS
from odps.accounts import BearerTokenAccount
logger.warning('Timeout met, killing the instance now.')
bearer_token = context().get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ['ODPS_PROJECT_NAME']
endpoint = os.environ['ODPS_RUNTIME_ENDPOINT']
o = ODPS(None,
None,
account=account,
project=project,
endpoint=endpoint)
o.stop_instance(os.environ['MARS_K8S_POD_NAMESPACE'])
else:
kv_store = context().kv_store()
kv_store[CUPID_LAST_IDLE_TIME_KEY] = str(self._last_active_time)
self.ref().check_instance_idle(_delay=10, _tell=True, _wait=False)
def execute(cls, ctx, op):
import pandas as pd
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
from cupid.io.table import CupidTableUploadSession
if op.is_terminal:
bearer_token = context().get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME', None)
odps_params = op.odps_params.copy()
if project:
odps_params['project'] = project
endpoint = os.environ.get(
'ODPS_RUNTIME_ENDPOINT') or odps_params['endpoint']
o = ODPS(None,
None,
account=account,
project=odps_params['project'],
endpoint=endpoint)
cupid_session = CupidSession(o)
project_name, table_name = op.table_name.split('.')
upload_session = CupidTableUploadSession(session=cupid_session,
table_name=table_name,
project_name=project_name,
handle=op.cupid_handle,
blocks=op.blocks)
upload_session.commit(overwrite=op.overwrite)
ctx[op.outputs[0].key] = pd.DataFrame()
def tile(cls, op):
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
bearer_token = context().get_bearer_token()
account = BearerTokenAccount(bearer_token)
o = ODPS(None, None, account=account, **op.odps_params)
cupid_session = CupidSession(o)
data_src = o.get_table(op.table_name)
logger.debug('Start creating upload session from cupid.')
upload_session = cupid_session.create_upload_session(data_src)
input_df = op.inputs[0]
out_chunks = []
out_chunk_shape = (0,) * len(input_df.shape)
blocks = {}
for chunk in input_df.chunks:
block_id = str(int(time.time())) + '_' + str(uuid.uuid4()).replace('-', '')
chunk_op = DataFrameWriteTableSplit(dtypes=op.dtypes, table_name=op.table_name,
partition_spec=op.partition_spec,
cupid_handle=to_str(upload_session.handle),
block_id=block_id, write_batch_size=op.write_batch_size)
out_chunk = chunk_op.new_chunk([chunk], shape=out_chunk_shape, index=chunk.index, dtypes=chunk.dtypes)
out_chunks.append(out_chunk)
blocks[block_id] = op.partition_spec
# build commit tree
combine_size = 8
chunks = out_chunks
while len(chunks) > combine_size:
new_chunks = []
for i in range(0, len(chunks), combine_size):
chks = chunks[i: i + combine_size]
if len(chks) == 1:
chk = chks[0]
else:
chk_op = DataFrameWriteTableCommit(dtypes=op.dtypes, is_terminal=False)
chk = chk_op.new_chunk(chks, shape=out_chunk_shape, dtypes=op.dtypes)
new_chunks.append(chk)
chunks = new_chunks
assert len(chunks) < combine_size
commit_table_op = DataFrameWriteTableCommit(dtypes=op.dtypes, table_name=op.table_name, blocks=blocks,
cupid_handle=to_str(upload_session.handle),
overwrite=op.overwrite, odps_params=op.odps_params,
is_terminal=True)
commit_table_chunk = commit_table_op.new_chunk(chunks, shape=out_chunk_shape, dtypes=op.dtypes)
out_df = op.outputs[0]
new_op = op.copy()
return new_op.new_dataframes(op.inputs, shape=out_df.shape,
dtypes=out_df.dtypes, chunks=[commit_table_chunk],
nsplits=((0,),) * len(out_chunk_shape))
def _handle_commit_table_upload_session(sock):
try:
cmd_len, = struct.unpack('<I', sock.recv(4))
# dict with odps_params, table_name, cupid_handle, blocks, overwrite
commit_config = pickle.loads(sock.recv(cmd_len))
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
from cupid.runtime import RuntimeContext
from cupid.io.table import CupidTableUploadSession
if not RuntimeContext.is_context_ready():
raise SystemError(
'No Mars cluster found, please create via `o.create_mars_cluster`.'
)
cupid_ctx = context()
odps_params = commit_config['odps_params']
bearer_token = cupid_ctx.get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME',
None) or odps_params['project']
endpoint = os.environ.get(
'ODPS_RUNTIME_ENDPOINT') or odps_params['endpoint']
o = ODPS(None,
None,
account=account,
project=project,
endpoint=endpoint)
cupid_session = CupidSession(o)
project_name, table_name = commit_config['table_name'].split('.')
upload_session = CupidTableUploadSession(
session=cupid_session,
table_name=table_name,
project_name=project_name,
handle=commit_config['cupid_handle'],
blocks=commit_config['blocks'])
upload_session.commit(overwrite=commit_config['overwrite'])
_write_request_result(sock)
except:
logger.exception('Failed to commit upload session')
_write_request_result(sock, False, exc_info=sys.exc_info())
def _handle_enum_table_partitions(sock):
try:
cmd_len, = struct.unpack('<I', sock.recv(4))
# dict with odps_params, table_name, partition
task_config = pickle.loads(sock.recv(cmd_len))
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import context
cupid_ctx = context()
odps_params = task_config['odps_params']
bearer_token = cupid_ctx.get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME',
None) or odps_params['project']
endpoint = os.environ.get(
'ODPS_RUNTIME_ENDPOINT') or odps_params['endpoint']
o = ODPS(None,
None,
account=account,
project=project,
endpoint=endpoint)
table = o.get_table(task_config['table_name'])
partition_desc = task_config.get('partition')
if not table.schema.partitions:
_write_request_result(sock, result=None)
elif partition_desc:
if check_partition_exist(table, partition_desc):
_write_request_result(sock, result=[partition_desc])
else:
parts = filter_partitions(o, list(table.partitions),
partition_desc)
_write_request_result(
sock, result=[str(pt.partition_spec) for pt in parts])
else:
_write_request_result(
sock,
result=[str(pt.partition_spec) for pt in table.partitions])
except:
logger.exception('Failed to create download session')
_write_request_result(sock, False, exc_info=sys.exc_info())
def _handle_create_table_upload_session(sock):
try:
cmd_len, = struct.unpack('<I', sock.recv(4))
# dict with odps_params, table_name
session_config = pickle.loads(sock.recv(cmd_len))
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
from cupid.runtime import RuntimeContext
if not RuntimeContext.is_context_ready():
raise SystemError(
'No Mars cluster found, please create via `o.create_mars_cluster`.'
)
cupid_ctx = context()
odps_params = session_config['odps_params']
bearer_token = cupid_ctx.get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME',
None) or odps_params['project']
endpoint = os.environ.get(
'ODPS_RUNTIME_ENDPOINT') or odps_params['endpoint']
o = ODPS(None,
None,
account=account,
project=project,
endpoint=endpoint)
cupid_session = CupidSession(o)
data_src = o.get_table(session_config['table_name'])
logger.debug('Start creating upload session from cupid.')
upload_session = cupid_session.create_upload_session(data_src)
ret_data = {
'handle': upload_session.handle,
}
_write_request_result(sock, result=ret_data)
except:
logger.exception('Failed to create upload session')
_write_request_result(sock, False, exc_info=sys.exc_info())
def execute(cls, ctx, op):
import pandas as pd
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
from cupid.io.table import CupidTableUploadSession
if op.is_terminal:
bearer_token = context().get_bearer_token()
account = BearerTokenAccount(bearer_token)
o = ODPS(None, None, account=account, **op.odps_params)
cupid_session = CupidSession(o)
project_name, table_name = op.table_name.split('.')
upload_session = CupidTableUploadSession(
session=cupid_session, table_name=table_name, project_name=project_name,
handle=op.cupid_handle, blocks=op.blocks)
upload_session.commit(overwrite=op.overwrite)
ctx[op.outputs[0].key] = pd.DataFrame()
def tile(cls, op):
import numpy as np
import pandas as pd
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
bearer_token = context().get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME', None)
odps_params = op.odps_params.copy()
if project:
odps_params['project'] = project
o = ODPS(None, None, account=account, **odps_params)
cupid_session = CupidSession(o)
df = op.outputs[0]
split_size = df.extra_params.chunk_bytes or CHUNK_LIMIT
data_src = o.get_table(op.table_name)
if op.partition is not None:
data_src = data_src.get_partition(op.partition)
logger.debug('Start creating download session from cupid.')
while True:
try:
download_session = cupid_session.create_download_session(
data_src, split_size=split_size, columns=op.columns)
break
except CupidError:
logger.debug(
'The number of splits exceeds 100000, split_size is {}'.
format(split_size))
if split_size >= MAX_CHUNK_SIZE:
raise
else:
split_size *= 2
logger.debug('%s table splits have been created.',
str(len(download_session.splits)))
out_chunks = []
# Ignore add_offset at this time.
op._add_offset = False
for idx, split in enumerate(download_session.splits):
chunk_op = DataFrameReadTableSplit(
cupid_handle=to_str(split.handle),
split_index=split.split_index,
split_file_start=split.split_file_start,
split_file_end=split.split_file_end,
schema_file_start=split.schema_file_start,
schema_file_end=split.schema_file_end,
add_offset=op.add_offset,
dtypes=op.dtypes,
sparse=op.sparse)
# the chunk shape is unknown
index_value = parse_index(pd.RangeIndex(0))
columns_value = parse_index(df.dtypes.index, store_data=True)
out_chunk = chunk_op.new_chunk(None,
shape=(np.nan, df.shape[1]),
dtypes=op.dtypes,
index_value=index_value,
columns_value=columns_value,
index=(idx, 0))
out_chunks.append(out_chunk)
if op.add_offset:
out_chunks = standardize_range_index(out_chunks)
new_op = op.copy()
nsplits = ((np.nan, ) * len(out_chunks), (df.shape[1], ))
return new_op.new_dataframes(None,
shape=df.shape,
dtypes=op.dtypes,
index_value=df.index_value,
columns_value=df.columns_value,
chunks=out_chunks,
nsplits=nsplits)
def _tile_cupid(cls, op):
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
from cupid.runtime import RuntimeContext
if not RuntimeContext.is_context_ready():
raise SystemError(
'No Mars cluster found, please create via `o.create_mars_cluster`.'
)
cupid_ctx = context()
bearer_token = cupid_ctx.get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME', None)
odps_params = op.odps_params.copy()
if project:
odps_params['project'] = project
endpoint = os.environ.get(
'ODPS_RUNTIME_ENDPOINT') or odps_params['endpoint']
o = ODPS(None,
None,
account=account,
project=odps_params['project'],
endpoint=endpoint)
cupid_session = CupidSession(o)
data_src = o.get_table(op.table_name)
logger.debug('Start creating upload session from cupid.')
upload_session = cupid_session.create_upload_session(data_src)
input_df = build_concatenated_rows_frame(op.inputs[0])
out_df = op.outputs[0]
out_chunks = []
out_chunk_shape = (0, ) * len(input_df.shape)
blocks = {}
for chunk in input_df.chunks:
block_id = str(int(time.time())) + '_' + str(uuid.uuid4()).replace(
'-', '')
chunk_op = DataFrameWriteTableSplit(
dtypes=op.dtypes,
table_name=op.table_name,
unknown_as_string=op.unknown_as_string,
partition_spec=op.partition_spec,
cupid_handle=to_str(upload_session.handle),
block_id=block_id,
write_batch_size=op.write_batch_size)
out_chunk = chunk_op.new_chunk([chunk],
shape=out_chunk_shape,
index=chunk.index,
index_value=out_df.index_value,
dtypes=chunk.dtypes)
out_chunks.append(out_chunk)
blocks[block_id] = op.partition_spec
# build commit tree
combine_size = 8
chunks = out_chunks
while len(chunks) >= combine_size:
new_chunks = []
for i in range(0, len(chunks), combine_size):
chks = chunks[i:i + combine_size]
if len(chks) == 1:
chk = chks[0]
else:
chk_op = DataFrameWriteTableCommit(dtypes=op.dtypes,
is_terminal=False)
chk = chk_op.new_chunk(chks,
shape=out_chunk_shape,
index_value=out_df.index_value,
dtypes=op.dtypes)
new_chunks.append(chk)
chunks = new_chunks
assert len(chunks) < combine_size
commit_table_op = DataFrameWriteTableCommit(dtypes=op.dtypes,
table_name=op.table_name,
blocks=blocks,
cupid_handle=to_str(
upload_session.handle),
overwrite=op.overwrite,
odps_params=op.odps_params,
is_terminal=True)
commit_table_chunk = commit_table_op.new_chunk(
chunks,
shape=out_chunk_shape,
dtypes=op.dtypes,
index_value=out_df.index_value)
new_op = op.copy()
return new_op.new_dataframes(op.inputs,
shape=out_df.shape,
index_value=out_df.index_value,
dtypes=out_df.dtypes,
columns_value=out_df.columns_value,
chunks=[commit_table_chunk],
nsplits=((0, ), ) * len(out_chunk_shape))
def _tile_cupid(cls, op):
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
from cupid.errors import CupidError
from mars.context import get_context
cupid_ctx = context()
bearer_token = cupid_ctx.get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME', None)
odps_params = op.odps_params.copy()
if project:
odps_params['project'] = project
endpoint = os.environ.get(
'ODPS_RUNTIME_ENDPOINT') or odps_params['endpoint']
o = ODPS(None,
None,
account=account,
project=odps_params['project'],
endpoint=endpoint)
cupid_session = CupidSession(o)
mars_context = get_context()
df = op.outputs[0]
split_size = df.extra_params.chunk_bytes or READ_CHUNK_LIMIT
out_dtypes = df.dtypes
out_shape = df.shape
out_columns_value = df.columns_value
if op.columns is not None:
out_dtypes = out_dtypes[op.columns]
out_shape = (df.shape[0], len(op.columns))
out_columns_value = parse_index(out_dtypes.index, store_data=True)
table_obj = o.get_table(op.table_name)
if not table_obj.schema.partitions:
data_srcs = [table_obj]
elif op.partition is not None and check_partition_exist(
table_obj, op.partition):
data_srcs = [table_obj.get_partition(op.partition)]
else:
data_srcs = list(table_obj.partitions)
if op.partition is not None:
data_srcs = filter_partitions(o, data_srcs, op.partition)
out_chunks = []
chunk_idx = 0
for data_src in data_srcs:
try:
data_store_size = data_src.size
except ODPSError:
# fail to get data size, just ignore
pass
else:
if data_store_size < split_size and mars_context is not None:
# get worker counts
worker_count = max(
len(mars_context.get_worker_addresses()), 1)
# data is too small, split as many as number of cores
split_size = data_store_size // worker_count
# at least 1M
split_size = max(split_size, 1 * 1024**2)
logger.debug(
'Input data size is too small, split_size is %s',
split_size)
logger.debug(
'Start creating download session of table %s from cupid, '
'columns: %s', op.table_name, op.columns)
while True:
try:
download_session = cupid_session.create_download_session(
data_src,
split_size=split_size,
columns=op.columns,
with_split_meta=op.with_split_meta_on_tile)
break
except CupidError:
logger.debug(
'The number of splits exceeds 100000, split_size is %s',
split_size)
if split_size >= MAX_CHUNK_SIZE:
raise
else:
split_size *= 2
logger.debug('%s table splits have been created.',
str(len(download_session.splits)))
meta_chunk_rows = [
split.meta_row_count for split in download_session.splits
]
if np.isnan(out_shape[0]):
est_chunk_rows = meta_chunk_rows
else:
sp_file_sizes = np.array([
sp.split_file_end - sp.split_file_start
for sp in download_session.splits
])
total_size = sp_file_sizes.sum()
ratio_chunk_rows = (sp_file_sizes * out_shape[0] //
total_size).tolist()
est_chunk_rows = [
mr if mr is not None else rr
for mr, rr in zip(meta_chunk_rows, ratio_chunk_rows)
]
partition_spec = str(data_src.partition_spec) \
if getattr(data_src, 'partition_spec', None) else None
logger.warning('Estimated chunk rows: %r', est_chunk_rows)
if len(download_session.splits) == 0:
logger.debug('Table %s has no data', op.table_name)
chunk_op = DataFrameReadTableSplit()
index_value = parse_index(pd.RangeIndex(0))
columns_value = parse_index(out_dtypes.index, store_data=True)
out_chunk = chunk_op.new_chunk(None,
shape=(np.nan, out_shape[1]),
dtypes=op.dtypes,
index_value=index_value,
columns_value=columns_value,
index=(chunk_idx, 0))
out_chunks.append(out_chunk)
chunk_idx += 1
else:
for idx, split in enumerate(download_session.splits):
chunk_op = DataFrameReadTableSplit(
cupid_handle=to_str(split.handle),
split_index=split.split_index,
split_file_start=split.split_file_start,
split_file_end=split.split_file_end,
schema_file_start=split.schema_file_start,
schema_file_end=split.schema_file_end,
add_offset=op.add_offset,
dtypes=out_dtypes,
sparse=op.sparse,
split_size=split_size,
string_as_binary=op.string_as_binary,
use_arrow_dtype=op.use_arrow_dtype,
estimate_rows=est_chunk_rows[idx],
partition_spec=partition_spec,
append_partitions=op.append_partitions,
meta_raw_size=split.meta_raw_size,
nrows=meta_chunk_rows[idx] or op.nrows,
memory_scale=op.memory_scale)
# the chunk shape is unknown
index_value = parse_index(pd.RangeIndex(0))
columns_value = parse_index(out_dtypes.index,
store_data=True)
out_chunk = chunk_op.new_chunk(None,
shape=(np.nan,
out_shape[1]),
dtypes=out_dtypes,
index_value=index_value,
columns_value=columns_value,
index=(chunk_idx, 0))
chunk_idx += 1
out_chunks.append(out_chunk)
if op.add_offset:
out_chunks = standardize_range_index(out_chunks)
new_op = op.copy()
nsplits = ((np.nan, ) * len(out_chunks), (out_shape[1], ))
return new_op.new_dataframes(None,
shape=out_shape,
dtypes=op.dtypes,
index_value=df.index_value,
columns_value=out_columns_value,
chunks=out_chunks,
nsplits=nsplits)
def tile(cls, op):
import numpy as np
import pandas as pd
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
bearer_token = context().get_bearer_token()
account = BearerTokenAccount(bearer_token)
o = ODPS(None, None, account=account, **op.odps_params)
cupid_session = CupidSession(o)
df = op.outputs[0]
split_size = df.extra_params.chunk_store_limit or options.tensor.chunk_store_limit
data_src = o.get_table(op.table_name)
if op.partition is not None:
data_src = data_src.get_partition(op.partition)
logger.debug('Start creating download session from cupid.')
download_session = cupid_session.create_download_session(
data_src, split_size=split_size)
logger.debug('%s table splits have been created.',
str(len(download_session.splits)))
out_chunks = []
out_count_chunks = []
for idx, split in enumerate(download_session.splits):
chunk_op = DataFrameReadTableSplit(
cupid_handle=to_str(split.handle),
split_index=split.split_index,
split_file_start=split.split_file_start,
split_file_end=split.split_file_end,
schema_file_start=split.schema_file_start,
schema_file_end=split.schema_file_end,
dtypes=op.dtypes,
sparse=op.sparse)
# the chunk shape is unknown
index_value = parse_index(pd.RangeIndex(0))
columns_value = parse_index(df.dtypes.index, store_data=True)
out_chunk, out_count_chunk = chunk_op.new_chunks(
None,
kws=[{
'shape': (np.nan, df.shape[1]),
'dtypes': op.dtypes,
'index_value': index_value,
'columns_value': columns_value,
'index': (idx, )
}, {
'shape': (1, ),
'index': (idx, )
}])
out_chunks.append(out_chunk)
out_count_chunks.append(out_count_chunk)
if op.add_offset:
output_chunks = []
for i, chunk in enumerate(out_chunks):
if i == 0:
output_chunks.append(chunk)
continue
counts = out_count_chunks[:i]
inputs = [chunk] + counts
output_chunk = DataFrameReadTableWithOffset(
dtypes=chunk.dtypes).new_chunk(
inputs,
shape=chunk.shape,
index=chunk.index,
dtypes=chunk.dtypes,
index_value=chunk.index_value,
columns_value=chunk.columns_value)
output_chunks.append(output_chunk)
else:
output_chunks = out_chunks
new_op = op.copy()
nsplits = ((np.nan, ) * len(output_chunks), (df.shape[1], ))
return new_op.new_dataframes(None,
shape=df.shape,
dtypes=op.dtypes,
index_value=df.index_value,
columns_value=df.columns_value,
chunks=output_chunks,
nsplits=nsplits)
def tile(cls, op):
import numpy as np
import pandas as pd
from odps import ODPS
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
from mars.context import get_context
cupid_ctx = context()
if cupid_ctx is None:
raise SystemError(
'No Mars cluster found, please create via `o.create_mars_cluster`.'
)
bearer_token = cupid_ctx.get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME', None)
odps_params = op.odps_params.copy()
if project:
odps_params['project'] = project
o = ODPS(None, None, account=account, **odps_params)
cupid_session = CupidSession(o)
mars_context = get_context()
df = op.outputs[0]
split_size = df.extra_params.chunk_bytes or READ_CHUNK_LIMIT
data_src = o.get_table(op.table_name)
if op.partition is not None:
data_src = data_src.get_partition(op.partition)
try:
data_store_size = data_src.size
except ODPSError:
# fail to get data size, just ignore
pass
else:
if data_store_size < split_size and mars_context is not None:
# get worker counts
worker_count = max(len(mars_context.get_worker_addresses()), 1)
# data is too small, split as many as number of cores
split_size = data_store_size // worker_count
# at least 1M
split_size = max(split_size, 1 * 1024**2)
logger.debug(
'Input data size is too small, split_size is {}'.format(
split_size))
logger.debug(
'Start creating download session of table {} from cupid.'.format(
op.table_name))
while True:
try:
download_session = cupid_session.create_download_session(
data_src, split_size=split_size, columns=op.columns)
break
except CupidError:
logger.debug(
'The number of splits exceeds 100000, split_size is {}'.
format(split_size))
if split_size >= MAX_CHUNK_SIZE:
raise
else:
split_size *= 2
logger.debug('%s table splits have been created.',
str(len(download_session.splits)))
if np.isnan(df.shape[0]):
est_chunk_rows = [None] * len(download_session.splits)
else:
sp_file_sizes = np.array([
sp.split_file_end - sp.split_file_start
for sp in download_session.splits
])
total_size = sp_file_sizes.sum()
est_chunk_rows = sp_file_sizes * df.shape[0] // total_size
logger.warning('Estimated chunk rows: %r', est_chunk_rows)
out_chunks = []
# Ignore add_offset at this time.
op._add_offset = False
if len(download_session.splits) == 0:
logger.debug('Table {} has no data'.format(op.table_name))
chunk_op = DataFrameReadTableSplit()
index_value = parse_index(pd.RangeIndex(0))
columns_value = parse_index(df.dtypes.index, store_data=True)
out_chunk = chunk_op.new_chunk(None,
shape=(np.nan, df.shape[1]),
dtypes=op.dtypes,
index_value=index_value,
columns_value=columns_value,
index=(0, 0))
out_chunks = [out_chunk]
else:
for idx, split in enumerate(download_session.splits):
chunk_op = DataFrameReadTableSplit(
cupid_handle=to_str(split.handle),
split_index=split.split_index,
split_file_start=split.split_file_start,
split_file_end=split.split_file_end,
schema_file_start=split.schema_file_start,
schema_file_end=split.schema_file_end,
add_offset=op.add_offset,
dtypes=op.dtypes,
sparse=op.sparse,
split_size=split_size,
use_arrow_dtype=op.use_arrow_dtype,
estimate_rows=est_chunk_rows[idx])
# the chunk shape is unknown
index_value = parse_index(pd.RangeIndex(0))
columns_value = parse_index(df.dtypes.index, store_data=True)
out_chunk = chunk_op.new_chunk(None,
shape=(np.nan, df.shape[1]),
dtypes=op.dtypes,
index_value=index_value,
columns_value=columns_value,
index=(idx, 0))
out_chunks.append(out_chunk)
if op.add_offset:
out_chunks = standardize_range_index(out_chunks)
new_op = op.copy()
nsplits = ((np.nan, ) * len(out_chunks), (df.shape[1], ))
return new_op.new_dataframes(None,
shape=df.shape,
dtypes=op.dtypes,
index_value=df.index_value,
columns_value=df.columns_value,
chunks=out_chunks,
nsplits=nsplits)
def _handle_create_table_download_session(sock):
try:
cmd_len, = struct.unpack('<I', sock.recv(4))
# dict with odps_params, table_name, partition, columns, worker_count, split_size, max_chunk_num
session_config = pickle.loads(sock.recv(cmd_len))
from odps import ODPS
from odps.errors import ODPSError
from odps.accounts import BearerTokenAccount
from cupid import CupidSession, context
from cupid.errors import CupidError
from cupid.runtime import RuntimeContext
if not RuntimeContext.is_context_ready():
raise SystemError(
'No Mars cluster found, please create via `o.create_mars_cluster`.'
)
cupid_ctx = context()
odps_params = session_config['odps_params']
bearer_token = cupid_ctx.get_bearer_token()
account = BearerTokenAccount(bearer_token)
project = os.environ.get('ODPS_PROJECT_NAME',
None) or odps_params['project']
endpoint = os.environ.get(
'ODPS_RUNTIME_ENDPOINT') or odps_params['endpoint']
o = ODPS(None,
None,
account=account,
project=project,
endpoint=endpoint)
cupid_session = CupidSession(o)
split_size = session_config['split_size']
table_name = session_config['table_name']
data_src = o.get_table(table_name)
if session_config.get('partition') is not None:
data_src = data_src.get_partition(session_config['partition'])
try:
data_store_size = data_src.size
except ODPSError:
# fail to get data size, just ignore
pass
else:
worker_count = session_config['worker_count']
if data_store_size < split_size and worker_count is not None:
# data is too small, split as many as number of cores
split_size = data_store_size // worker_count
# at least 1M
split_size = max(split_size, 1 * 1024**2)
logger.debug(
'Input data size is too small, split_size is {}'.format(
split_size))
max_chunk_num = session_config['max_chunk_num']
columns = session_config['columns']
with_split_meta = session_config.get('with_split_meta_on_tile')
logger.debug(
'Start creating download session of table %s from cupid, columns %r',
table_name, columns)
while True:
try:
download_session = cupid_session.create_download_session(
data_src,
split_size=split_size,
columns=columns,
with_split_meta=with_split_meta)
break
except CupidError:
logger.debug(
'The number of splits exceeds 100000, split_size is {}'.
format(split_size))
if split_size >= max_chunk_num:
raise
else:
split_size *= 2
ret_data = {
'splits': download_session.splits,
'split_size': split_size,
}
_write_request_result(sock, result=ret_data)
except:
logger.exception('Failed to create download session')
_write_request_result(sock, False, exc_info=sys.exc_info())
