def test_encode(): """ Test encoding of a message into 'RecordIO' format. """ try: encoder = recordio.Encoder(lambda s: bytes(json.dumps(s), "UTF-8")) except Exception as exception: raise MesosException( "Error instantiating 'RecordIO' encoder: {error}".format( error=exception)) try: message = { "type": "ATTACH_CONTAINER_OUTPUT", "containerId": "123456789" } encoded = encoder.encode(message) except Exception as exception: raise MesosException( "Error encoding 'RecordIO' message: {error}".format( error=exception)) string = json.dumps(message) assert encoded == bytes(str(len(string)) + "\n" + string, "UTF-8")
def __init__(self, agent_url, parent_container_id, user=None, cmd=None, args=None, interactive=False, tty=False): # Store relevant parameters of the call for later. self.cmd = cmd self.interactive = interactive self.tty = tty self.args = args # Get the URL to the agent running the task. self.agent_url = urllib.parse.urljoin(agent_url, 'api/v1') # Grab a reference to the container ID for the task. self.parent_id = parent_container_id self.user = user # Generate a new UUID for the nested container # used to run commands passed to `task exec`. self.container_id = str(uuid.uuid4()) # Set up a recordio encoder and decoder # for any incoming and outgoing messages. self.encoder = recordio.Encoder( lambda s: bytes(json.dumps(s, ensure_ascii=False), "UTF-8")) self.decoder = recordio.Decoder( lambda s: json.loads(s.decode("UTF-8"))) # Set up queues to send messages between threads used for # reading/writing to STDIN/STDOUT/STDERR and threads # sending/receiving data over the network. self.input_queue = Queue() self.output_queue = Queue() # Set up an event to block attaching # input until attaching output is complete. self.attach_input_event = threading.Event() self.attach_input_event.clear() # Set up an event to block printing the output # until an attach input event has successfully # been established. self.print_output_event = threading.Event() self.print_output_event.clear() # Set up an event to block the main thread # from exiting until signaled to do so. self.exit_event = threading.Event() self.exit_event.clear() # Use a class variable to store exceptions thrown on # other threads and raise them on the main thread before # exiting. self.exception = None
def test_encode_decode(): """ Test encoding/decoding of a message and records into 'RecordIO' format. """ total_messages = 10 try: encoder = recordio.Encoder(lambda s: bytes(json.dumps(s), "UTF-8")) except Exception as exception: raise MesosException( "Error instantiating 'RecordIO' encoder: {error}".format( error=exception)) try: decoder = recordio.Decoder(lambda s: json.loads(s.decode("UTF-8"))) except Exception as exception: raise MesosException( "Error instantiating 'RecordIO' decoder: {error}".format( error=exception)) try: message = { "type": "ATTACH_CONTAINER_OUTPUT", "containerId": "123456789" } encoded = b"" for _ in range(total_messages): encoded += encoder.encode(message) except Exception as exception: raise MesosException( "Error encoding 'RecordIO' message: {error}".format( error=exception)) try: all_records = [] offset = 0 chunk_size = 5 while offset < len(encoded): records = decoder.decode(encoded[offset:offset + chunk_size]) all_records.extend(records) offset += chunk_size assert len(all_records) == total_messages for record in all_records: assert record == message except Exception as exception: raise MesosException( "Error decoding 'RecordIO' messages: {error}".format( error=exception))
def __init__(self, mesos_master_url, task_id, cmd=None, args=None, interactive=False, tty=False): # Store relevant parameters of the call for later. self.cmd = cmd self.interactive = interactive self.tty = tty self.args = args self._mesos_master_url = mesos_master_url master = Master(get_master_state(self._mesos_master_url)) # Get the task and make sure its container was launched by the UCR. # Since task's containers are launched by the UCR by default, we want # to allow most tasks to pass through unchecked. The only exception is # when a task has an explicit container specified and it is not of type # "MESOS". Having a type of "MESOS" implies that it was launched by the # UCR -- all other types imply it was not. task_obj = master.task(task_id) if "container" in task_obj.dict(): if "type" in task_obj.dict()["container"]: if task_obj.dict()["container"]["type"] != "MESOS": raise MesosException( "This command is only supported for tasks" " launched by the Universal Container Runtime (UCR).") # Get the URL to the agent running the task. self.agent_url = urllib.parse.urljoin(task_obj.slave().http_url(), 'api/v1') # Grab a reference to the container ID for the task. self.parent_id = master.get_container_id(task_id) if "user" in task_obj.dict(): self.user = task_obj.dict()['user'] else: self.user = None # Generate a new UUID for the nested container # used to run commands passed to `task exec`. self.container_id = str(uuid.uuid4()) # Set up a recordio encoder and decoder # for any incoming and outgoing messages. self.encoder = recordio.Encoder( lambda s: bytes(json.dumps(s, ensure_ascii=False), "UTF-8")) self.decoder = recordio.Decoder( lambda s: json.loads(s.decode("UTF-8"))) # Set up queues to send messages between threads used for # reading/writing to STDIN/STDOUT/STDERR and threads # sending/receiving data over the network. self.input_queue = Queue() self.output_queue = Queue() # Set up an event to block attaching # input until attaching output is complete. self.attach_input_event = threading.Event() self.attach_input_event.clear() # Set up an event to block printing the output # until an attach input event has successfully # been established. self.print_output_event = threading.Event() self.print_output_event.clear() # Set up an event to block the main thread # from exiting until signaled to do so. self.exit_event = threading.Event() self.exit_event.clear() # Use a class variable to store exceptions thrown on # other threads and raise them on the main thread before # exiting. self.exception = None
def __init__(self, master, config, task_id): # Get the task and make sure its container was launched by the UCR. # Since task's containers are launched by the UCR by default, we want # to allow most tasks to pass through unchecked. The only exception is # when a task has an explicit container specified and it is not of type # "MESOS". Having a type of "MESOS" implies that it was launched by the # UCR -- all other types imply it was not. try: tasks = get_tasks(master, config, query={'task_id': task_id}) except Exception as exception: raise CLIException("Unable to get task with ID {task_id}" " from leading master '{master}': {error}" .format(task_id=task_id, master=master, error=exception)) running_tasks = [t for t in tasks if t["state"] == "TASK_RUNNING"] matching_tasks = [t for t in running_tasks if t["id"] == task_id] if not matching_tasks: raise CLIException("Unable to find running task '{task_id}'" " from leading master '{master}'" .format(task_id=task_id, master=master)) if len(matching_tasks) > 1: raise CLIException("More than one task matching id '{id}'" .format(id=task_id)) task_obj = matching_tasks[0] if "container" in task_obj: if "type" in task_obj["container"]: if task_obj["container"]["type"] != "MESOS": raise CLIException( "This command is only supported for tasks" " launched by the Universal Container Runtime (UCR).") # Get the scheme of the agent scheme = "https://" if config.agent_ssl() else "http://" # Get the URL to the agent running the task. agent_addr = util.sanitize_address( scheme + get_agent_address(task_obj["slave_id"], master, config)) self.agent_url = mesos.http.simple_urljoin(agent_addr, "api/v1") # Get the agent's task path by checking the `state` endpoint. try: self.container_id = get_container_id(task_obj) except CLIException as exception: raise CLIException("Could not get container ID of task '{id}'" " from agent '{addr}': {error}" .format(id=task_id, addr=agent_addr, error=exception)) # Set up a recordio encoder and decoder # for any incoming and outgoing messages. self.encoder = recordio.Encoder( lambda s: bytes(json.dumps(s, ensure_ascii=False), "UTF-8")) self.decoder = recordio.Decoder( lambda s: json.loads(s.decode("UTF-8"))) # Set up queues to send messages between threads used for # reading/writing to STDIN/STDOUT/STDERR and threads # sending/receiving data over the network. self.input_queue = Queue() self.output_queue = Queue() # Set up an event to block attaching # input until attaching output is complete. self.attach_input_event = threading.Event() self.attach_input_event.clear() # Set up an event to block printing the output # until an attach input event has successfully # been established. self.print_output_event = threading.Event() self.print_output_event.clear() # Set up an event to block the main thread # from exiting until signaled to do so. self.exit_event = threading.Event() self.exit_event.clear() # Use a class variable to store exceptions thrown on # other threads and raise them on the main thread before # exiting. self.exception = None # Default values for the TaskIO. self.cmd = None self.args = None self.interactive = False self.tty = False self.output_thread_entry_point = None self.config = config # Allow an exit sequence to be used to break the CLIs attachment to # the remote task. Depending on the call, this may be disabled, or # the exit sequence to be used may be overwritten. self.supports_exit_sequence = False self.exit_sequence = b'\x10\x11' # Ctrl-p, Ctrl-q self.exit_sequence_detected = False