def test(endpoint_id=None, tasks=10, hostname=None, port=None): tasks_rq = RedisQueue(f'task_{endpoint_id}', hostname) results_rq = RedisQueue(f'results', hostname) fxs = FuncXSerializer() ser_code = fxs.serialize(slow_double) fn_code = fxs.pack_buffers([ser_code]) while True: try: x = results_rq.get(timeout=1) except: print("No more results left") break tasks_rq.connect() results_rq.connect() start = time.time() for i in range(tasks): ser_args = fxs.serialize([i]) ser_kwargs = fxs.serialize({'duration':0}) input_data = fxs.pack_buffers([ser_args, ser_kwargs]) payload = fn_code + input_data tasks_rq.put(f"0{i}", payload) for i in range(tasks): res = results_rq.get(timeout=1) print("Result : ", res) delta = time.time() - start print("Time to complete {} tasks: {:8.3f} s".format(tasks, delta)) print("Throughput : {:8.3f} Tasks/s".format(tasks / delta)) return delta
def dont_run_yet(endpoint_id=None, tasks=10, duration=1, hostname=None): # tasks_rq = EndpointQueue(f'task_{endpoint_id}', hostname) tasks_channel = RedisPubSub(hostname) tasks_channel.connect() redis_client = tasks_channel.redis_client redis_client.ping() fxs = FuncXSerializer() ser_code = fxs.serialize(slow_double) fn_code = fxs.pack_buffers([ser_code]) start = time.time() task_ids = {} for i in range(tasks): time.sleep(duration) task_id = str(uuid.uuid4()) print("Task_id : ", task_id) ser_args = fxs.serialize([i]) ser_kwargs = fxs.serialize({"duration": duration}) input_data = fxs.pack_buffers([ser_args, ser_kwargs]) payload = fn_code + input_data container_id = "RAW" task = Task(redis_client, task_id, container_id, serializer="", payload=payload) task.endpoint = endpoint_id task.status = TaskState.WAITING_FOR_EP # tasks_rq.enqueue(task) tasks_channel.put(endpoint_id, task) task_ids[i] = task_id d1 = time.time() - start print(f"Time to launch {tasks} tasks: {d1:8.3f} s") delay = 5 print(f"Sleeping {delay} seconds") time.sleep(delay) print(f"Launched {tasks} tasks") for i in range(tasks): task_id = task_ids[i] print("Task_id : ", task_id) task = Task.from_id(redis_client, task_id) # TODO: wait for task result... time.sleep(duration) try: result = fxs.deserialize(task.result) print(f"Result : {result}") except Exception as e: print(f"Task failed with exception:{e}") pass delta = time.time() - start print(f"Time to complete {tasks} tasks: {delta:8.3f} s") print(f"Throughput : {tasks / delta:8.3f} Tasks/s") return delta
def serialize_fx_inputs(*args, **kwargs): from funcx.serialize import FuncXSerializer fx_serializer = FuncXSerializer() ser_args = fx_serializer.serialize(args) ser_kwargs = fx_serializer.serialize(kwargs) payload = fx_serializer.pack_buffers([ser_args, ser_kwargs]) return payload
def serialize_fx_inputs(*args, **kwargs): """Pack and serialize inputs """ fx_serializer = FuncXSerializer() ser_args = fx_serializer.serialize(args) ser_kwargs = fx_serializer.serialize(kwargs) payload = fx_serializer.pack_buffers([ser_args, ser_kwargs]) return payload
class Batch: """Utility class for creating batch submission in funcX""" def __init__(self): self.tasks = [] self.fx_serializer = FuncXSerializer() def add(self, *args, endpoint_id=None, function_id=None, **kwargs): """Add an function invocation to a batch submission Parameters ---------- *args : Any Args as specified by the function signature endpoint_id : uuid str Endpoint UUID string. Required function_id : uuid str Function UUID string. Required asynchronous : bool Whether or not to run the function asynchronously Returns ------- None """ assert endpoint_id is not None, "endpoint_id key-word argument must be set" assert function_id is not None, "function_id key-word argument must be set" ser_args = self.fx_serializer.serialize(args) ser_kwargs = self.fx_serializer.serialize(kwargs) payload = self.fx_serializer.pack_buffers([ser_args, ser_kwargs]) data = { 'endpoint': endpoint_id, 'function': function_id, 'payload': payload } self.tasks.append(data) def prepare(self): """Prepare the payloads to be post to web service in a batch Parameters ---------- Returns ------- payloads in dictionary, Dict[str, list] """ data = {'tasks': []} for task in self.tasks: new_task = (task['function'], task['endpoint'], task['payload']) data['tasks'].append(new_task) return data
def dont_run_yet(endpoint_id=None, tasks=10, duration=1, hostname=None): tasks_rq = EndpointQueue(f"task_{endpoint_id}", hostname) fxs = FuncXSerializer() ser_code = fxs.serialize(slow_double) fn_code = fxs.pack_buffers([ser_code]) tasks_rq.connect() start = time.time() task_ids = {} for i in range(tasks): task_id = str(uuid.uuid4()) ser_args = fxs.serialize([i]) ser_kwargs = fxs.serialize({"duration": duration}) input_data = fxs.pack_buffers([ser_args, ser_kwargs]) payload = fn_code + input_data container_id = "RAW" task = Task(tasks_rq.redis_client, task_id, container_id, serializer="", payload=payload) tasks_rq.enqueue(task) task_ids[i] = task_id d1 = time.time() - start print(f"Time to launch {tasks} tasks: {d1:8.3f} s") print(f"Launched {tasks} tasks") for i in range(tasks): task_id = task_ids[i] task = Task.from_id(tasks_rq.redis_client, task_id) # TODO: wait for task result... time.sleep(2) print(f"Result: {task.result}") # res = results_rq.get('result', timeout=300) # print("Result : ", res) delta = time.time() - start print(f"Time to complete {tasks} tasks: {delta:8.3f} s") print(f"Throughput : {tasks / delta:8.3f} Tasks/s") return delta
def test(endpoint_id=None, tasks=10, duration=1, hostname=None, port=None): tasks_rq = RedisQueue(f'task_{endpoint_id}', hostname) results_rq = RedisQueue('results', hostname) fxs = FuncXSerializer() ser_code = fxs.serialize(slow_double) fn_code = fxs.pack_buffers([ser_code]) tasks_rq.connect() results_rq.connect() while True: try: _ = results_rq.get(timeout=1) except Exception: print("No more results left") break start = time.time() for i in range(tasks): ser_args = fxs.serialize([i]) ser_kwargs = fxs.serialize({'duration': duration}) input_data = fxs.pack_buffers([ser_args, ser_kwargs]) payload = fn_code + input_data container_id = "odd" if i % 2 else "even" tasks_rq.put(f"0{i};{container_id}", payload) d1 = time.time() - start print("Time to launch {} tasks: {:8.3f} s".format(tasks, d1)) print(f"Launched {tasks} tasks") for i in range(tasks): _ = results_rq.get(timeout=300) # print("Result : ", res) delta = time.time() - start print("Time to complete {} tasks: {:8.3f} s".format(tasks, delta)) print("Throughput : {:8.3f} Tasks/s".format(tasks / delta)) return delta
results_queue = Queue() # set_file_logger('executor.log', name='funcx_endpoint', level=logging.DEBUG) htex = HighThroughputExecutor(interchange_local=True, passthrough=True) htex.start(results_passthrough=results_queue) htex._start_remote_interchange_process() fx_serializer = FuncXSerializer() for i in range(10): task_id = str(uuid.uuid4()) args = (i, ) kwargs = {} fn_code = fx_serializer.serialize(double) ser_code = fx_serializer.pack_buffers([fn_code]) ser_params = fx_serializer.pack_buffers( [fx_serializer.serialize(args), fx_serializer.serialize(kwargs)]) payload = Task(task_id, "RAW", ser_code + ser_params) f = htex.submit_raw(payload.pack()) time.sleep(0.5) for i in range(10): result_package = results_queue.get() # print("Result package : ", result_package) r = pickle.loads(result_package) result = fx_serializer.deserialize(r["result"]) print(f"Result:{i}: {result}")
print("Started the endpoint {}".format(endpoint_name)) print("Wating 10 seconds for the endpoint to start") time.sleep(10) # Connect to the task and result redis queue endpoint_id = args.endpoint_id tasks_rq = RedisQueue(f'task_{endpoint_id}', args.redis_hostname) results_rq = RedisQueue(f'results', args.redis_hostname) tasks_rq.connect() results_rq.connect() print("Redis queue connected") # Create an instance of funcx serializer and serialize the function fxs = FuncXSerializer() ser_code = fxs.serialize(sleep) fn_code = fxs.pack_buffers([ser_code]) print("Code serialized") # Define the test function def test(tasks=10, durs=[5, 10, 20], timeout=None): # Make sure there is no previous result left while True: try: x = results_rq.get(timeout=1) except: print("No more results left") break start_submit = time.time() time_table = {} rounds = 3 #counts = [[1, 5, 20], [5, 20, 1], [20, 1, 5]]
class FuncXClient(throttling.ThrottledBaseClient): """Main class for interacting with the funcX service Holds helper operations for performing common tasks with the funcX service. """ TOKEN_DIR = os.path.expanduser("~/.funcx/credentials") TOKEN_FILENAME = 'funcx_sdk_tokens.json' CLIENT_ID = '4cf29807-cf21-49ec-9443-ff9a3fb9f81c' def __init__(self, http_timeout=None, funcx_home=os.path.join('~', '.funcx'), force_login=False, fx_authorizer=None, funcx_service_address='https://api.funcx.org/v1', **kwargs): """ Initialize the client Parameters ---------- http_timeout: int Timeout for any call to service in seconds. Default is no timeout force_login: bool Whether to force a login to get new credentials. fx_authorizer:class:`GlobusAuthorizer <globus_sdk.authorizers.base.GlobusAuthorizer>`: A custom authorizer instance to communicate with funcX. Default: ``None``, will be created. funcx_service_address: str The address of the funcX web service to communicate with. Default: https://api.funcx.org/v1 Keyword arguments are the same as for BaseClient. """ self.func_table = {} self.ep_registration_path = 'register_endpoint_2' self.funcx_home = os.path.expanduser(funcx_home) if not os.path.exists(self.TOKEN_DIR): os.makedirs(self.TOKEN_DIR) tokens_filename = os.path.join(self.TOKEN_DIR, self.TOKEN_FILENAME) self.native_client = NativeClient( client_id=self.CLIENT_ID, app_name="FuncX SDK", token_storage=JSONTokenStorage(tokens_filename)) # TODO: if fx_authorizer is given, we still need to get an authorizer for Search fx_scope = "https://auth.globus.org/scopes/facd7ccc-c5f4-42aa-916b-a0e270e2c2a9/all" search_scope = "urn:globus:auth:scope:search.api.globus.org:all" scopes = [fx_scope, search_scope, "openid"] search_authorizer = None if not fx_authorizer: self.native_client.login( requested_scopes=scopes, no_local_server=kwargs.get("no_local_server", True), no_browser=kwargs.get("no_browser", True), refresh_tokens=kwargs.get("refresh_tokens", True), force=force_login) all_authorizers = self.native_client.get_authorizers_by_scope( requested_scopes=scopes) fx_authorizer = all_authorizers[fx_scope] search_authorizer = all_authorizers[search_scope] openid_authorizer = all_authorizers["openid"] super(FuncXClient, self).__init__("funcX", environment='funcx', authorizer=fx_authorizer, http_timeout=http_timeout, base_url=funcx_service_address, **kwargs) self.fx_serializer = FuncXSerializer() authclient = AuthClient(authorizer=openid_authorizer) user_info = authclient.oauth2_userinfo() self.searcher = SearchHelper(authorizer=search_authorizer, owner_uuid=user_info['sub']) self.funcx_service_address = funcx_service_address def version_check(self): """Check this client version meets the service's minimum supported version. """ resp = self.get("version", params={"service": "all"}) versions = resp.data if "min_ep_version" not in versions: raise VersionMismatch( "Failed to retrieve version information from funcX service.") min_ep_version = versions['min_ep_version'] if ENDPOINT_VERSION is None: raise VersionMismatch( "You do not have the funcx endpoint installed. You can use 'pip install funcx-endpoint'." ) if ENDPOINT_VERSION < min_ep_version: raise VersionMismatch( f"Your version={ENDPOINT_VERSION} is lower than the " f"minimum version for an endpoint: {min_ep_version}. Please update." ) def logout(self): """Remove credentials from your local system """ self.native_client.logout() def update_table(self, return_msg, task_id): """ Parses the return message from the service and updates the internal func_tables Parameters ---------- return_msg : str Return message received from the funcx service task_id : str task id string """ if isinstance(return_msg, str): r_dict = json.loads(return_msg) else: r_dict = return_msg status = {'pending': True} if 'result' in r_dict: try: r_obj = self.fx_serializer.deserialize(r_dict['result']) completion_t = r_dict['completion_t'] except Exception: raise SerializationError("Result Object Deserialization") else: status.update({ 'pending': False, 'result': r_obj, 'completion_t': completion_t }) self.func_table[task_id] = status elif 'exception' in r_dict: try: r_exception = self.fx_serializer.deserialize( r_dict['exception']) completion_t = r_dict['completion_t'] logger.info(f"Exception : {r_exception}") except Exception: raise SerializationError( "Task's exception object deserialization") else: status.update({ 'pending': False, 'exception': r_exception, 'completion_t': completion_t }) self.func_table[task_id] = status return status def get_task(self, task_id): """Get a funcX task. Parameters ---------- task_id : str UUID of the task Returns ------- dict Task block containing "status" key. """ if task_id in self.func_table: return self.func_table[task_id] r = self.get("tasks/{task_id}".format(task_id=task_id)) logger.debug("Response string : {}".format(r)) try: rets = self.update_table(r.text, task_id) except Exception as e: raise e return rets def get_result(self, task_id): """ Get the result of a funcX task Parameters ---------- task_id: str UUID of the task Returns ------- Result obj: If task completed Raises ------ Exception obj: Exception due to which the task failed """ task = self.get_task(task_id) if task['pending'] is True: raise Exception("Task pending") else: if 'result' in task: return task['result'] else: logger.warning("We have an exception : {}".format( task['exception'])) task['exception'].reraise() def get_batch_status(self, task_id_list): """ Request status for a batch of task_ids """ assert isinstance(task_id_list, list), "get_batch_status expects a list of task ids" pending_task_ids = [ t for t in task_id_list if t not in self.func_table ] results = {} if pending_task_ids: payload = {'task_ids': pending_task_ids} r = self.post("/batch_status", json_body=payload) logger.debug("Response string : {}".format(r)) pending_task_ids = set(pending_task_ids) for task_id in task_id_list: if task_id in pending_task_ids: try: data = r['results'][task_id] rets = self.update_table(data, task_id) results[task_id] = rets except KeyError: logger.debug( "Task {} info was not available in the batch status") except Exception: logger.exception( "Failure while unpacking results fom get_batch_status") else: results[task_id] = self.func_table[task_id] return results def get_batch_result(self, task_id_list): """ Request results for a batch of task_ids """ pass def run(self, *args, endpoint_id=None, function_id=None, **kwargs): """Initiate an invocation Parameters ---------- *args : Any Args as specified by the function signature endpoint_id : uuid str Endpoint UUID string. Required function_id : uuid str Function UUID string. Required asynchronous : bool Whether or not to run the function asynchronously Returns ------- task_id : str UUID string that identifies the task """ assert endpoint_id is not None, "endpoint_id key-word argument must be set" assert function_id is not None, "function_id key-word argument must be set" batch = self.create_batch() batch.add(*args, endpoint_id=endpoint_id, function_id=function_id, **kwargs) r = self.batch_run(batch) """ Create a future to deal with the result funcx_future = FuncXFuture(self, task_id, async_poll) if not asynchronous: return funcx_future.result() # Return the result return funcx_future """ return r[0] def create_batch(self): """ Create a Batch instance to handle batch submission in funcX Parameters ---------- Returns ------- Batch instance Status block containing "status" key. """ batch = Batch() return batch def batch_run(self, batch): """Initiate a batch of tasks to funcX Parameters ---------- batch: a Batch object Returns ------- task_ids : a list of UUID strings that identify the tasks """ servable_path = 'submit' assert isinstance(batch, Batch), "Requires a Batch object as input" assert len(batch.tasks) > 0, "Requires a non-empty batch" data = batch.prepare() # Send the data to funcX r = self.post(servable_path, json_body=data) if r.http_status != 200: raise HTTPError(r) if r.get("status", "Failure") == "Failure": raise MalformedResponse("FuncX Request failed: {}".format( r.get("reason", "Unknown"))) return r['task_uuids'] def map_run(self, *args, endpoint_id=None, function_id=None, asynchronous=False, **kwargs): """Initiate an invocation Parameters ---------- *args : Any Args as specified by the function signature endpoint_id : uuid str Endpoint UUID string. Required function_id : uuid str Function UUID string. Required asynchronous : bool Whether or not to run the function asynchronously Returns ------- task_id : str UUID string that identifies the task """ servable_path = 'submit_batch' assert endpoint_id is not None, "endpoint_id key-word argument must be set" assert function_id is not None, "function_id key-word argument must be set" ser_kwargs = self.fx_serializer.serialize(kwargs) batch_payload = [] iterator = args[0] for arg in iterator: ser_args = self.fx_serializer.serialize((arg, )) payload = self.fx_serializer.pack_buffers([ser_args, ser_kwargs]) batch_payload.append(payload) data = { 'endpoints': [endpoint_id], 'func': function_id, 'payload': batch_payload, 'is_async': asynchronous } # Send the data to funcX r = self.post(servable_path, json_body=data) if r.http_status != 200: raise Exception(r) if r.get("status", "Failure") == "Failure": raise MalformedResponse("FuncX Request failed: {}".format( r.get("reason", "Unknown"))) return r['task_uuids'] def register_endpoint(self, name, endpoint_uuid, metadata=None, endpoint_version=None): """Register an endpoint with the funcX service. Parameters ---------- name : str Name of the endpoint endpoint_uuid : str The uuid of the endpoint metadata : dict endpoint metadata, see default_config example endpoint_version: str Version string to be passed to the webService as a compatibility check Returns ------- A dict {'endopoint_id' : <>, 'address' : <>, 'client_ports': <>} """ self.version_check() data = { "endpoint_name": name, "endpoint_uuid": endpoint_uuid, "version": endpoint_version } if metadata: data['meta'] = metadata r = self.post(self.ep_registration_path, json_body=data) if r.http_status != 200: raise HTTPError(r) # Return the result return r.data def get_containers(self, name, description=None): """Register a DLHub endpoint with the funcX service and get the containers to launch. Parameters ---------- name : str Name of the endpoint description : str Description of the endpoint Returns ------- int The port to connect to and a list of containers """ registration_path = 'get_containers' data = {"endpoint_name": name, "description": description} r = self.post(registration_path, json_body=data) if r.http_status != 200: raise HTTPError(r) # Return the result return r.data['endpoint_uuid'], r.data['endpoint_containers'] def get_container(self, container_uuid, container_type): """Get the details of a container for staging it locally. Parameters ---------- container_uuid : str UUID of the container in question container_type : str The type of containers that will be used (Singularity, Shifter, Docker) Returns ------- dict The details of the containers to deploy """ container_path = f'containers/{container_uuid}/{container_type}' r = self.get(container_path) if r.http_status != 200: raise HTTPError(r) # Return the result return r.data['container'] def get_endpoint_status(self, endpoint_uuid): """Get the status reports for an endpoint. Parameters ---------- endpoint_uuid : str UUID of the endpoint in question Returns ------- dict The details of the endpoint's stats """ stats_path = f'endpoints/{endpoint_uuid}/status' r = self.get(stats_path) if r.http_status != 200: raise HTTPError(r) # Return the result return r.data def register_function(self, function, function_name=None, container_uuid=None, description=None, public=False, group=None, searchable=True): """Register a function code with the funcX service. Parameters ---------- function : Python Function The function to be registered for remote execution function_name : str The entry point (function name) of the function. Default: None container_uuid : str Container UUID from registration with funcX description : str Description of the file public : bool Whether or not the function is publicly accessible. Default = False group : str A globus group uuid to share this function with searchable : bool If true, the function will be indexed into globus search with the appropriate permissions Returns ------- function uuid : str UUID identifier for the registered function """ registration_path = 'register_function' source_code = "" try: source_code = getsource(function) except OSError: logger.error( "Failed to find source code during function registration.") serialized_fn = self.fx_serializer.serialize(function) packed_code = self.fx_serializer.pack_buffers([serialized_fn]) data = { "function_name": function.__name__, "function_code": packed_code, "function_source": source_code, "container_uuid": container_uuid, "entry_point": function_name if function_name else function.__name__, "description": description, "public": public, "group": group, "searchable": searchable } logger.info("Registering function : {}".format(data)) r = self.post(registration_path, json_body=data) if r.http_status != 200: raise HTTPError(r) func_uuid = r.data['function_uuid'] # Return the result return func_uuid def update_function(self, func_uuid, function): pass def search_function(self, q, offset=0, limit=10, advanced=False): """Search for function via the funcX service Parameters ---------- q : str free-form query string offset : int offset into total results limit : int max number of results to return advanced : bool allows elastic-search like syntax in query string Returns ------- FunctionSearchResults """ return self.searcher.search_function(q, offset=offset, limit=limit, advanced=advanced) def search_endpoint(self, q, scope='all', owner_id=None): """ Parameters ---------- q scope : str Can be one of {'all', 'my-endpoints', 'shared-with-me'} owner_id should be urn like f"urn:globus:auth:identity:{owner_uuid}" Returns ------- """ return self.searcher.search_endpoint(q, scope=scope, owner_id=owner_id) def register_container(self, location, container_type, name='', description=''): """Register a container with the funcX service. Parameters ---------- location : str The location of the container (e.g., its docker url). Required container_type : str The type of containers that will be used (Singularity, Shifter, Docker). Required name : str A name for the container. Default = '' description : str A description to associate with the container. Default = '' Returns ------- str The id of the container """ container_path = 'containers' payload = { 'name': name, 'location': location, 'description': description, 'type': container_type } r = self.post(container_path, json_body=payload) if r.http_status != 200: raise HTTPError(r) # Return the result return r.data['container_id'] def add_to_whitelist(self, endpoint_id, function_ids): """Adds the function to the endpoint's whitelist Parameters ---------- endpoint_id : str The uuid of the endpoint function_ids : list A list of function id's to be whitelisted Returns ------- json The response of the request """ req_path = f'endpoints/{endpoint_id}/whitelist' if not isinstance(function_ids, list): function_ids = [function_ids] payload = {'func': function_ids} r = self.post(req_path, json_body=payload) if r.http_status != 200: raise HTTPError(r) # Return the result return r def get_whitelist(self, endpoint_id): """List the endpoint's whitelist Parameters ---------- endpoint_id : str The uuid of the endpoint Returns ------- json The response of the request """ req_path = f'endpoints/{endpoint_id}/whitelist' r = self.get(req_path) if r.http_status != 200: raise HTTPError(r) # Return the result return r def delete_from_whitelist(self, endpoint_id, function_ids): """List the endpoint's whitelist Parameters ---------- endpoint_id : str The uuid of the endpoint function_ids : list A list of function id's to be whitelisted Returns ------- json The response of the request """ if not isinstance(function_ids, list): function_ids = [function_ids] res = [] for fid in function_ids: req_path = f'endpoints/{endpoint_id}/whitelist/{fid}' r = self.delete(req_path) if r.http_status != 200: raise HTTPError(r) res.append(r) # Return the result return res
class FuncXClient(BaseClient): """Main class for interacting with the funcX service Holds helper operations for performing common tasks with the funcX service. """ TOKEN_DIR = os.path.expanduser("~/.funcx/credentials") CLIENT_ID = '4cf29807-cf21-49ec-9443-ff9a3fb9f81c' def __init__(self, http_timeout=None, funcx_home=os.path.join('~', '.funcx'), force_login=False, fx_authorizer=None, funcx_service_address='https://dev.funcx.org/api/v1', **kwargs): """ Initialize the client Parameters ---------- http_timeout: int Timeout for any call to service in seconds. Default is no timeout force_login: bool Whether to force a login to get new credentials. fx_authorizer:class:`GlobusAuthorizer <globus_sdk.authorizers.base.GlobusAuthorizer>`: A custom authorizer instance to communicate with funcX. Default: ``None``, will be created. service_address: str The address of the funcX web service to communicate with. Default: https://dev.funcx.org/api/v1 Keyword arguments are the same as for BaseClient. """ self.ep_registration_path = 'register_endpoint_2' self.funcx_home = os.path.expanduser(funcx_home) native_client = NativeClient(client_id=self.CLIENT_ID) fx_scope = "https://auth.globus.org/scopes/facd7ccc-c5f4-42aa-916b-a0e270e2c2a9/all" if not fx_authorizer: native_client.login( requested_scopes=[fx_scope], no_local_server=kwargs.get("no_local_server", True), no_browser=kwargs.get("no_browser", True), refresh_tokens=kwargs.get("refresh_tokens", True), force=force_login) all_authorizers = native_client.get_authorizers_by_scope( requested_scopes=[fx_scope]) fx_authorizer = all_authorizers[fx_scope] super(FuncXClient, self).__init__("funcX", environment='funcx', authorizer=fx_authorizer, http_timeout=http_timeout, base_url=funcx_service_address, **kwargs) self.fx_serializer = FuncXSerializer() def logout(self): """Remove credentials from your local system """ logout() def get_task_status(self, task_id): """Get the status of a funcX task. Parameters ---------- task_id : str UUID of the task Returns ------- dict Status block containing "status" key. """ r = self.get("{task_id}/status".format(task_id=task_id)) return json.loads(r.text) def get_result(self, task_id): """ Get the result of a funcX task Parameters ---------- task_id: str UUID of the task Returns ------- Result obj: If task completed Raises ------ Exception obj: Exception due to which the task failed """ r = self.get("{task_id}/status".format(task_id=task_id)) logger.info(f"Got from globus : {r}") r_dict = json.loads(r.text) if 'result' in r_dict: try: r_obj = self.fx_serializer.deserialize(r_dict['result']) except Exception: raise Exception( "Failure during deserialization of the result object") else: return r_obj elif 'exception' in r_dict: try: r_exception = self.fx_serializer.deserialize( r_dict['exception']) logger.info(f"Exception : {r_exception}") except Exception: raise Exception( "Failure during deserialization of the Task's exception object" ) else: r_exception.reraise() else: raise Exception("Task pending") def run(self, *args, endpoint_id=None, function_id=None, asynchronous=False, **kwargs): """Initiate an invocation Parameters ---------- *args : Any Args as specified by the function signature endpoint_id : uuid str Endpoint UUID string. Required function_id : uuid str Function UUID string. Required asynchronous : bool Whether or not to run the function asynchronously Returns ------- task_id : str UUID string that identifies the task """ servable_path = 'submit' assert endpoint_id is not None, "endpoint_id key-word argument must be set" assert function_id is not None, "function_id key-word argument must be set" ser_args = self.fx_serializer.serialize(args) ser_kwargs = self.fx_serializer.serialize(kwargs) payload = self.fx_serializer.pack_buffers([ser_args, ser_kwargs]) data = { 'endpoint': endpoint_id, 'func': function_id, 'payload': payload, 'is_async': asynchronous } # Send the data to funcX r = self.post(servable_path, json_body=data) if r.http_status is not 200: raise Exception(r) if 'task_uuid' not in r: raise MalformedResponse(r) """ Create a future to deal with the result funcx_future = FuncXFuture(self, task_id, async_poll) if not asynchronous: return funcx_future.result() # Return the result return funcx_future """ return r['task_uuid'] def register_endpoint(self, name, endpoint_uuid, description=None): """Register an endpoint with the funcX service. Parameters ---------- name : str Name of the endpoint endpoint_uuid : str The uuid of the endpoint description : str Description of the endpoint Returns ------- A dict {'endopoint_id' : <>, 'address' : <>, 'client_ports': <>} """ data = { "endpoint_name": name, "endpoint_uuid": endpoint_uuid, "description": description } r = self.post(self.ep_registration_path, json_body=data) if r.http_status is not 200: raise Exception(r) # Return the result return r.data def get_containers(self, name, description=None): """Register a DLHub endpoint with the funcX service and get the containers to launch. Parameters ---------- name : str Name of the endpoint description : str Description of the endpoint Returns ------- int The port to connect to and a list of containers """ registration_path = 'get_containers' data = {"endpoint_name": name, "description": description} r = self.post(registration_path, json_body=data) if r.http_status is not 200: raise Exception(r) # Return the result return r.data['endpoint_uuid'], r.data['endpoint_containers'] def get_container(self, container_uuid, container_type): """Get the details of a container for staging it locally. Parameters ---------- container_uuid : str UUID of the container in question container_type : str The type of containers that will be used (Singularity, Shifter, Docker) Returns ------- dict The details of the containers to deploy """ container_path = f'containers/{container_uuid}/{container_type}' r = self.get(container_path) if r.http_status is not 200: raise Exception(r) # Return the result return r.data['container'] def register_function(self, function, function_name=None, container_uuid=None, description=None): """Register a function code with the funcX service. Parameters ---------- function : Python Function The function to be registered for remote execution function_name : str The entry point (function name) of the function. Default: None container_uuid : str Container UUID from registration with funcX description : str Description of the file Returns ------- function uuid : str UUID identifier for the registered function """ registration_path = 'register_function' serialized_fn = self.fx_serializer.serialize(function) packed_code = self.fx_serializer.pack_buffers([serialized_fn]) data = { "function_name": function.__name__, "function_code": packed_code, "container_uuid": container_uuid, "entry_point": function_name if function_name else function.__name__, "description": description } logger.info("Registering function : {}".format(data)) r = self.post(registration_path, json_body=data) if r.http_status is not 200: raise Exception(r) # Return the result return r.data['function_uuid']
class FuncXClient: """Main class for interacting with the funcX service Holds helper operations for performing common tasks with the funcX service. """ TOKEN_DIR = os.path.expanduser("~/.funcx/credentials") TOKEN_FILENAME = "funcx_sdk_tokens.json" FUNCX_SDK_CLIENT_ID = os.environ.get( "FUNCX_SDK_CLIENT_ID", "4cf29807-cf21-49ec-9443-ff9a3fb9f81c" ) FUNCX_SCOPE = os.environ.get( "FUNCX_SCOPE", "https://auth.globus.org/scopes/facd7ccc-c5f4-42aa-916b-a0e270e2c2a9/all", ) def __init__( self, http_timeout=None, funcx_home=_FUNCX_HOME, force_login=False, fx_authorizer=None, search_authorizer=None, openid_authorizer=None, funcx_service_address=None, check_endpoint_version=False, asynchronous=False, loop=None, results_ws_uri=None, use_offprocess_checker=True, environment=None, **kwargs, ): """ Initialize the client Parameters ---------- http_timeout: int Timeout for any call to service in seconds. Default is no timeout force_login: bool Whether to force a login to get new credentials. fx_authorizer:class:`GlobusAuthorizer \ <globus_sdk.authorizers.base.GlobusAuthorizer>`: A custom authorizer instance to communicate with funcX. Default: ``None``, will be created. search_authorizer:class:`GlobusAuthorizer \ <globus_sdk.authorizers.base.GlobusAuthorizer>`: A custom authorizer instance to communicate with Globus Search. Default: ``None``, will be created. openid_authorizer:class:`GlobusAuthorizer \ <globus_sdk.authorizers.base.GlobusAuthorizer>`: A custom authorizer instance to communicate with OpenID. Default: ``None``, will be created. funcx_service_address: str For internal use only. The address of the web service. results_ws_uri: str For internal use only. The address of the websocket service. environment: str For internal use only. The name of the environment to use. asynchronous: bool Should the API use asynchronous interactions with the web service? Currently only impacts the run method Default: False loop: AbstractEventLoop If asynchronous mode is requested, then you can provide an optional event loop instance. If None, then we will access asyncio.get_event_loop() Default: None use_offprocess_checker: Bool, Use this option to disable the offprocess_checker in the FuncXSerializer used by the client. Default: True Keyword arguments are the same as for BaseClient. """ # resolve URLs if not set if funcx_service_address is None: funcx_service_address = get_web_service_url(environment) if results_ws_uri is None: results_ws_uri = get_web_socket_url(environment) self.func_table = {} self.use_offprocess_checker = use_offprocess_checker self.funcx_home = os.path.expanduser(funcx_home) self.session_task_group_id = str(uuid.uuid4()) if not os.path.exists(self.TOKEN_DIR): os.makedirs(self.TOKEN_DIR) tokens_filename = os.path.join(self.TOKEN_DIR, self.TOKEN_FILENAME) self.native_client = NativeClient( client_id=self.FUNCX_SDK_CLIENT_ID, app_name="FuncX SDK", token_storage=JSONTokenStorage(tokens_filename), ) # TODO: if fx_authorizer is given, we still need to get an authorizer for Search search_scope = "urn:globus:auth:scope:search.api.globus.org:all" scopes = [self.FUNCX_SCOPE, search_scope, "openid"] if not fx_authorizer or not search_authorizer or not openid_authorizer: self.native_client.login( requested_scopes=scopes, no_local_server=kwargs.get("no_local_server", True), no_browser=kwargs.get("no_browser", True), refresh_tokens=kwargs.get("refresh_tokens", True), force=force_login, ) all_authorizers = self.native_client.get_authorizers_by_scope( requested_scopes=scopes ) fx_authorizer = all_authorizers[self.FUNCX_SCOPE] search_authorizer = all_authorizers[search_scope] openid_authorizer = all_authorizers["openid"] self.web_client = FuncxWebClient( base_url=funcx_service_address, authorizer=fx_authorizer ) self.fx_serializer = FuncXSerializer( use_offprocess_checker=self.use_offprocess_checker ) authclient = AuthClient(authorizer=openid_authorizer) user_info = authclient.oauth2_userinfo() self.searcher = SearchHelper( authorizer=search_authorizer, owner_uuid=user_info["sub"] ) self.funcx_service_address = funcx_service_address self.check_endpoint_version = check_endpoint_version self.version_check() self.results_ws_uri = results_ws_uri self.asynchronous = asynchronous if asynchronous: self.loop = loop if loop else asyncio.get_event_loop() # Start up an asynchronous polling loop in the background self.ws_polling_task = WebSocketPollingTask( self, self.loop, init_task_group_id=self.session_task_group_id, results_ws_uri=self.results_ws_uri, ) else: self.loop = None def version_check(self): """Check this client version meets the service's minimum supported version.""" resp = self.web_client.get_version() versions = resp.data if "min_ep_version" not in versions: raise VersionMismatch( "Failed to retrieve version information from funcX service." ) min_ep_version = versions["min_ep_version"] min_sdk_version = versions["min_sdk_version"] if self.check_endpoint_version: if ENDPOINT_VERSION is None: raise VersionMismatch( "You do not have the funcx endpoint installed. " "You can use 'pip install funcx-endpoint'." ) if LooseVersion(ENDPOINT_VERSION) < LooseVersion(min_ep_version): raise VersionMismatch( f"Your version={ENDPOINT_VERSION} is lower than the " f"minimum version for an endpoint: {min_ep_version}. " "Please update. " f"pip install funcx-endpoint>={min_ep_version}" ) else: if LooseVersion(SDK_VERSION) < LooseVersion(min_sdk_version): raise VersionMismatch( f"Your version={SDK_VERSION} is lower than the " f"minimum version for funcx SDK: {min_sdk_version}. " "Please update. " f"pip install funcx>={min_sdk_version}" ) def logout(self): """Remove credentials from your local system""" self.native_client.logout() def update_table(self, return_msg, task_id): """Parses the return message from the service and updates the internal func_table Parameters ---------- return_msg : str Return message received from the funcx service task_id : str task id string """ if isinstance(return_msg, str): r_dict = json.loads(return_msg) else: r_dict = return_msg r_status = r_dict.get("status", "unknown") status = {"pending": True, "status": r_status} if "result" in r_dict: try: r_obj = self.fx_serializer.deserialize(r_dict["result"]) completion_t = r_dict["completion_t"] except Exception: raise SerializationError("Result Object Deserialization") else: status.update( {"pending": False, "result": r_obj, "completion_t": completion_t} ) self.func_table[task_id] = status elif "exception" in r_dict: try: r_exception = self.fx_serializer.deserialize(r_dict["exception"]) completion_t = r_dict["completion_t"] logger.info(f"Exception : {r_exception}") except Exception: raise SerializationError("Task's exception object deserialization") else: status.update( { "pending": False, "exception": r_exception, "completion_t": completion_t, } ) self.func_table[task_id] = status return status def get_task(self, task_id): """Get a funcX task. Parameters ---------- task_id : str UUID of the task Returns ------- dict Task block containing "status" key. """ if task_id in self.func_table: return self.func_table[task_id] r = self.web_client.get_task(task_id) logger.debug(f"Response string : {r}") try: rets = self.update_table(r.text, task_id) except Exception as e: raise e return rets def get_result(self, task_id): """Get the result of a funcX task Parameters ---------- task_id: str UUID of the task Returns ------- Result obj: If task completed Raises ------ Exception obj: Exception due to which the task failed """ task = self.get_task(task_id) if task["pending"] is True: raise TaskPending(task["status"]) else: if "result" in task: return task["result"] else: logger.warning("We have an exception : {}".format(task["exception"])) task["exception"].reraise() def get_batch_result(self, task_id_list): """Request status for a batch of task_ids""" assert isinstance( task_id_list, list ), "get_batch_result expects a list of task ids" pending_task_ids = [t for t in task_id_list if t not in self.func_table] results = {} if pending_task_ids: r = self.web_client.get_batch_status(pending_task_ids) logger.debug(f"Response string : {r}") pending_task_ids = set(pending_task_ids) for task_id in task_id_list: if task_id in pending_task_ids: try: data = r["results"][task_id] rets = self.update_table(data, task_id) results[task_id] = rets except KeyError: logger.debug("Task {} info was not available in the batch status") except Exception: logger.exception( "Failure while unpacking results fom get_batch_result" ) else: results[task_id] = self.func_table[task_id] return results def run(self, *args, endpoint_id=None, function_id=None, **kwargs): """Initiate an invocation Parameters ---------- *args : Any Args as specified by the function signature endpoint_id : uuid str Endpoint UUID string. Required function_id : uuid str Function UUID string. Required asynchronous : bool Whether or not to run the function asynchronously Returns ------- task_id : str UUID string that identifies the task if asynchronous is False funcX Task: asyncio.Task A future that will eventually resolve into the function's result if asynchronous is True """ assert endpoint_id is not None, "endpoint_id key-word argument must be set" assert function_id is not None, "function_id key-word argument must be set" batch = self.create_batch() batch.add(*args, endpoint_id=endpoint_id, function_id=function_id, **kwargs) r = self.batch_run(batch) return r[0] def create_batch(self, task_group_id=None): """ Create a Batch instance to handle batch submission in funcX Parameters ---------- task_group_id : str Override the session wide session_task_group_id with a different task_group_id for this batch. If task_group_id is not specified, it will default to using the client's session_task_group_id Returns ------- Batch instance Status block containing "status" key. """ if not task_group_id: task_group_id = self.session_task_group_id batch = Batch(task_group_id=task_group_id) return batch def batch_run(self, batch): """Initiate a batch of tasks to funcX Parameters ---------- batch: a Batch object Returns ------- task_ids : a list of UUID strings that identify the tasks """ assert isinstance(batch, Batch), "Requires a Batch object as input" assert len(batch.tasks) > 0, "Requires a non-empty batch" data = batch.prepare() # Send the data to funcX r = self.web_client.submit(data) task_uuids = [] for result in r["results"]: task_id = result["task_uuid"] task_uuids.append(task_id) if result["http_status_code"] != 200: # this method of handling errors for a batch response is not # ideal, as it will raise any error in the multi-response, # but it will do until batch_run is deprecated in favor of Executer handle_response_errors(result) if self.asynchronous: task_group_id = r["task_group_id"] asyncio_tasks = [] for task_id in task_uuids: funcx_task = FuncXTask(task_id) asyncio_task = self.loop.create_task(funcx_task.get_result()) asyncio_tasks.append(asyncio_task) self.ws_polling_task.add_task(funcx_task) self.ws_polling_task.put_task_group_id(task_group_id) return asyncio_tasks return task_uuids def map_run( self, *args, endpoint_id=None, function_id=None, asynchronous=False, **kwargs ): """Initiate an invocation Parameters ---------- *args : Any Args as specified by the function signature endpoint_id : uuid str Endpoint UUID string. Required function_id : uuid str Function UUID string. Required asynchronous : bool Whether or not to run the function asynchronously Returns ------- task_id : str UUID string that identifies the task """ assert endpoint_id is not None, "endpoint_id key-word argument must be set" assert function_id is not None, "function_id key-word argument must be set" ser_kwargs = self.fx_serializer.serialize(kwargs) batch_payload = [] iterator = args[0] for arg in iterator: ser_args = self.fx_serializer.serialize((arg,)) payload = self.fx_serializer.pack_buffers([ser_args, ser_kwargs]) batch_payload.append(payload) data = { "endpoints": [endpoint_id], "func": function_id, "payload": batch_payload, "is_async": asynchronous, } # Send the data to funcX r = self.web_client.submit_batch(data) return r["task_uuids"] def register_endpoint( self, name, endpoint_uuid, metadata=None, endpoint_version=None ): """Register an endpoint with the funcX service. Parameters ---------- name : str Name of the endpoint endpoint_uuid : str The uuid of the endpoint metadata : dict endpoint metadata, see default_config example endpoint_version: str Version string to be passed to the webService as a compatibility check Returns ------- A dict {'endpoint_id' : <>, 'address' : <>, 'client_ports': <>} """ self.version_check() r = self.web_client.register_endpoint( endpoint_name=name, endpoint_id=endpoint_uuid, metadata=metadata, endpoint_version=endpoint_version, ) return r.data def get_containers(self, name, description=None): """Register a DLHub endpoint with the funcX service and get the containers to launch. Parameters ---------- name : str Name of the endpoint description : str Description of the endpoint Returns ------- int The port to connect to and a list of containers """ data = {"endpoint_name": name, "description": description} r = self.web_client.post("get_containers", data=data) return r.data["endpoint_uuid"], r.data["endpoint_containers"] def get_container(self, container_uuid, container_type): """Get the details of a container for staging it locally. Parameters ---------- container_uuid : str UUID of the container in question container_type : str The type of containers that will be used (Singularity, Shifter, Docker) Returns ------- dict The details of the containers to deploy """ self.version_check() r = self.web_client.get(f"containers/{container_uuid}/{container_type}") return r.data["container"] def get_endpoint_status(self, endpoint_uuid): """Get the status reports for an endpoint. Parameters ---------- endpoint_uuid : str UUID of the endpoint in question Returns ------- dict The details of the endpoint's stats """ r = self.web_client.get_endpoint_status(endpoint_uuid) return r.data def register_function( self, function, function_name=None, container_uuid=None, description=None, public=False, group=None, searchable=True, ): """Register a function code with the funcX service. Parameters ---------- function : Python Function The function to be registered for remote execution function_name : str The entry point (function name) of the function. Default: None container_uuid : str Container UUID from registration with funcX description : str Description of the file public : bool Whether or not the function is publicly accessible. Default = False group : str A globus group uuid to share this function with searchable : bool If true, the function will be indexed into globus search with the appropriate permissions Returns ------- function uuid : str UUID identifier for the registered function """ data = FunctionRegistrationData( function=function, failover_source="", container_uuid=container_uuid, entry_point=function_name, description=description, public=public, group=group, searchable=searchable, serializer=self.fx_serializer, ) logger.info(f"Registering function : {data}") r = self.web_client.register_function(data) return r.data["function_uuid"] def search_function(self, q, offset=0, limit=10, advanced=False): """Search for function via the funcX service Parameters ---------- q : str free-form query string offset : int offset into total results limit : int max number of results to return advanced : bool allows elastic-search like syntax in query string Returns ------- FunctionSearchResults """ return self.searcher.search_function( q, offset=offset, limit=limit, advanced=advanced ) def search_endpoint(self, q, scope="all", owner_id=None): """ Parameters ---------- q scope : str Can be one of {'all', 'my-endpoints', 'shared-with-me'} owner_id should be urn like f"urn:globus:auth:identity:{owner_uuid}" Returns ------- """ return self.searcher.search_endpoint(q, scope=scope, owner_id=owner_id) def register_container(self, location, container_type, name="", description=""): """Register a container with the funcX service. Parameters ---------- location : str The location of the container (e.g., its docker url). Required container_type : str The type of containers that will be used (Singularity, Shifter, Docker). Required name : str A name for the container. Default = '' description : str A description to associate with the container. Default = '' Returns ------- str The id of the container """ payload = { "name": name, "location": location, "description": description, "type": container_type, } r = self.web_client.post("containers", data=payload) return r.data["container_id"] def add_to_whitelist(self, endpoint_id, function_ids): """Adds the function to the endpoint's whitelist Parameters ---------- endpoint_id : str The uuid of the endpoint function_ids : list A list of function id's to be whitelisted Returns ------- json The response of the request """ return self.web_client.whitelist_add(endpoint_id, function_ids) def get_whitelist(self, endpoint_id): """List the endpoint's whitelist Parameters ---------- endpoint_id : str The uuid of the endpoint Returns ------- json The response of the request """ return self.web_client.get_whitelist(endpoint_id) def delete_from_whitelist(self, endpoint_id, function_ids): """List the endpoint's whitelist Parameters ---------- endpoint_id : str The uuid of the endpoint function_ids : list A list of function id's to be whitelisted Returns ------- json The response of the request """ if not isinstance(function_ids, list): function_ids = [function_ids] res = [] for fid in function_ids: res.append(self.web_client.whitelist_remove(endpoint_id, fid)) return res
class FuncXClient(throttling.ThrottledBaseClient): """Main class for interacting with the funcX service Holds helper operations for performing common tasks with the funcX service. """ TOKEN_DIR = os.path.expanduser("~/.funcx/credentials") TOKEN_FILENAME = 'funcx_sdk_tokens.json' CLIENT_ID = '4cf29807-cf21-49ec-9443-ff9a3fb9f81c' def __init__(self, http_timeout=None, funcx_home=os.path.join('~', '.funcx'), force_login=False, fx_authorizer=None, funcx_service_address='https://funcx.org/api/v1', **kwargs): """ Initialize the client Parameters ---------- http_timeout: int Timeout for any call to service in seconds. Default is no timeout force_login: bool Whether to force a login to get new credentials. fx_authorizer:class:`GlobusAuthorizer <globus_sdk.authorizers.base.GlobusAuthorizer>`: A custom authorizer instance to communicate with funcX. Default: ``None``, will be created. service_address: str The address of the funcX web service to communicate with. Default: https://dev.funcx.org/api/v1 Keyword arguments are the same as for BaseClient. """ self.func_table = {} self.ep_registration_path = 'register_endpoint_2' self.funcx_home = os.path.expanduser(funcx_home) if not os.path.exists(self.TOKEN_DIR): os.makedirs(self.TOKEN_DIR) tokens_filename = os.path.join(self.TOKEN_DIR, self.TOKEN_FILENAME) self.native_client = NativeClient( client_id=self.CLIENT_ID, app_name="FuncX SDK", token_storage=JSONTokenStorage(tokens_filename)) fx_scope = "https://auth.globus.org/scopes/facd7ccc-c5f4-42aa-916b-a0e270e2c2a9/all" if not fx_authorizer: self.native_client.login( requested_scopes=[fx_scope], no_local_server=kwargs.get("no_local_server", True), no_browser=kwargs.get("no_browser", True), refresh_tokens=kwargs.get("refresh_tokens", True), force=force_login) all_authorizers = self.native_client.get_authorizers_by_scope( requested_scopes=[fx_scope]) fx_authorizer = all_authorizers[fx_scope] super(FuncXClient, self).__init__("funcX", environment='funcx', authorizer=fx_authorizer, http_timeout=http_timeout, base_url=funcx_service_address, **kwargs) self.fx_serializer = FuncXSerializer() def logout(self): """Remove credentials from your local system """ self.native_client.logout() def update_table(self, return_msg, task_id): """ Parses the return message from the service and updates the internal func_tables Parameters ---------- return_msg : str Return message received from the funcx service task_id : str task id string """ if isinstance(return_msg, str): r_dict = json.loads(return_msg) else: r_dict = return_msg status = {'pending': True} if 'result' in r_dict: try: r_obj = self.fx_serializer.deserialize(r_dict['result']) except Exception: raise Exception( "Failure during deserialization of the result object") else: status.update({'pending': 'False', 'result': r_obj}) self.func_table[task_id] = status elif 'exception' in r_dict: try: r_exception = self.fx_serializer.deserialize( r_dict['exception']) logger.info(f"Exception : {r_exception}") except Exception: raise Exception( "Failure during deserialization of the Task's exception object" ) else: status.update({'pending': 'False', 'exception': r_exception}) self.func_table[task_id] = status return status def get_task_status(self, task_id): """Get the status of a funcX task. Parameters ---------- task_id : str UUID of the task Returns ------- dict Status block containing "status" key. """ if task_id in self.func_table: return self.func_table[task_id] r = self.get("{task_id}/status".format(task_id=task_id)) logger.debug("Response string : {}".format(r)) try: rets = self.update_table(r.text, task_id) except Exception as e: raise e return rets def get_result(self, task_id): """ Get the result of a funcX task Parameters ---------- task_id: str UUID of the task Returns ------- Result obj: If task completed Raises ------ Exception obj: Exception due to which the task failed """ status = self.get_task_status(task_id) if status['pending'] is True: raise Exception("Task pending") else: if 'result' in status: return status['result'] else: logger.warn("We have an exception : {}".format( status['exception'])) status['exception'].reraise() def get_batch_status(self, task_id_list): """ Request status for a batch of task_ids """ assert isinstance(task_id_list, list), "get_batch_status expects a list of task ids" pending_task_ids = [ t for t in task_id_list if t not in self.func_table ] results = {} if pending_task_ids: payload = {'task_ids': pending_task_ids} r = self.post("/batch_status", json_body=payload) logger.debug("Response string : {}".format(r)) pending_task_ids = set(pending_task_ids) for task_id in task_id_list: if task_id in pending_task_ids: try: data = r['results'][task_id] rets = self.update_table(data, task_id) results[task_id] = rets except KeyError: logger.debug( "Task {} info was not available in the batch status") except Exception as e: logger.exception( "Failure while unpacking results fom get_batch_status") else: results[task_id] = self.func_table[task_id] return results def get_batch_result(self, task_id_list): """ Request results for a batch of task_ids """ pass def run(self, *args, endpoint_id=None, function_id=None, asynchronous=False, **kwargs): """Initiate an invocation Parameters ---------- *args : Any Args as specified by the function signature endpoint_id : uuid str Endpoint UUID string. Required function_id : uuid str Function UUID string. Required asynchronous : bool Whether or not to run the function asynchronously Returns ------- task_id : str UUID string that identifies the task """ servable_path = 'submit' assert endpoint_id is not None, "endpoint_id key-word argument must be set" assert function_id is not None, "function_id key-word argument must be set" ser_args = self.fx_serializer.serialize(args) ser_kwargs = self.fx_serializer.serialize(kwargs) payload = self.fx_serializer.pack_buffers([ser_args, ser_kwargs]) data = { 'endpoint': endpoint_id, 'func': function_id, 'payload': payload, 'is_async': asynchronous } # Send the data to funcX r = self.post(servable_path, json_body=data) if r.http_status is not 200: raise Exception(r) if 'task_uuid' not in r: raise MalformedResponse(r) """ Create a future to deal with the result funcx_future = FuncXFuture(self, task_id, async_poll) if not asynchronous: return funcx_future.result() # Return the result return funcx_future """ return r['task_uuid'] def map_run(self, *args, endpoint_id=None, function_id=None, asynchronous=False, **kwargs): """Initiate an invocation Parameters ---------- *args : Any Args as specified by the function signature endpoint_id : uuid str Endpoint UUID string. Required function_id : uuid str Function UUID string. Required asynchronous : bool Whether or not to run the function asynchronously Returns ------- task_id : str UUID string that identifies the task """ servable_path = 'submit_batch' assert endpoint_id is not None, "endpoint_id key-word argument must be set" assert function_id is not None, "function_id key-word argument must be set" ser_kwargs = self.fx_serializer.serialize(kwargs) batch_payload = [] iterator = args[0] for arg in iterator: ser_args = self.fx_serializer.serialize((arg, )) payload = self.fx_serializer.pack_buffers([ser_args, ser_kwargs]) batch_payload.append(payload) data = { 'endpoints': [endpoint_id], 'func': function_id, 'payload': batch_payload, 'is_async': asynchronous } # Send the data to funcX r = self.post(servable_path, json_body=data) if r.http_status is not 200: raise Exception(r) if 'task_uuids' not in r: raise MalformedResponse(r) return r['task_uuids'] def register_endpoint(self, name, endpoint_uuid, description=None): """Register an endpoint with the funcX service. Parameters ---------- name : str Name of the endpoint endpoint_uuid : str The uuid of the endpoint description : str Description of the endpoint Returns ------- A dict {'endopoint_id' : <>, 'address' : <>, 'client_ports': <>} """ data = { "endpoint_name": name, "endpoint_uuid": endpoint_uuid, "description": description } r = self.post(self.ep_registration_path, json_body=data) if r.http_status is not 200: raise Exception(r) # Return the result return r.data def get_containers(self, name, description=None): """Register a DLHub endpoint with the funcX service and get the containers to launch. Parameters ---------- name : str Name of the endpoint description : str Description of the endpoint Returns ------- int The port to connect to and a list of containers """ registration_path = 'get_containers' data = {"endpoint_name": name, "description": description} r = self.post(registration_path, json_body=data) if r.http_status is not 200: raise Exception(r) # Return the result return r.data['endpoint_uuid'], r.data['endpoint_containers'] def get_container(self, container_uuid, container_type): """Get the details of a container for staging it locally. Parameters ---------- container_uuid : str UUID of the container in question container_type : str The type of containers that will be used (Singularity, Shifter, Docker) Returns ------- dict The details of the containers to deploy """ container_path = f'containers/{container_uuid}/{container_type}' r = self.get(container_path) if r.http_status is not 200: raise Exception(r) # Return the result return r.data['container'] def register_function(self, function, function_name=None, container_uuid=None, description=None): """Register a function code with the funcX service. Parameters ---------- function : Python Function The function to be registered for remote execution function_name : str The entry point (function name) of the function. Default: None container_uuid : str Container UUID from registration with funcX description : str Description of the file Returns ------- function uuid : str UUID identifier for the registered function """ registration_path = 'register_function' serialized_fn = self.fx_serializer.serialize(function) packed_code = self.fx_serializer.pack_buffers([serialized_fn]) data = { "function_name": function.__name__, "function_code": packed_code, "container_uuid": container_uuid, "entry_point": function_name if function_name else function.__name__, "description": description } logger.info("Registering function : {}".format(data)) r = self.post(registration_path, json_body=data) if r.http_status is not 200: raise Exception(r) # Return the result return r.data['function_uuid'] def register_container(self, location, container_type, name='', description=''): """Register a container with the funcX service. Parameters ---------- location : str The location of the container (e.g., its docker url). Required container_type : str The type of containers that will be used (Singularity, Shifter, Docker). Required name : str A name for the container. Default = '' description : str A description to associate with the container. Default = '' Returns ------- str The id of the container """ container_path = f'containers' payload = { 'name': name, 'location': location, 'description': description, 'type': container_type } r = self.post(container_path, json_body=payload) if r.http_status is not 200: raise Exception(r) # Return the result return r.data['container_id']