def worker(sock): """ Called by a worker process after the fork(). """ signal.signal(SIGHUP, SIG_DFL) signal.signal(SIGCHLD, SIG_DFL) signal.signal(SIGTERM, SIG_DFL) # restore the handler for SIGINT, # it's useful for debugging (show the stacktrace before exit) signal.signal(SIGINT, signal.default_int_handler) # Read the socket using fdopen instead of socket.makefile() because the latter # seems to be very slow; note that we need to dup() the file descriptor because # otherwise writes also cause a seek that makes us miss data on the read side. infile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) outfile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) exit_code = 0 try: worker_main(infile, outfile) except SystemExit as exc: exit_code = compute_real_exit_code(exc.code) finally: try: outfile.flush() except Exception: pass return exit_code
def worker(listen_sock): # Redirect stdout to stderr os.dup2(2, 1) sys.stdout = sys.stderr # The sys.stdout object is different from file descriptor 1 # Manager sends SIGHUP to request termination of workers in the pool def handle_sighup(*args): assert should_exit() signal.signal(SIGHUP, handle_sighup) # Cleanup zombie children def handle_sigchld(*args): pid = status = None try: while (pid, status) != (0, 0): pid, status = os.waitpid(0, os.WNOHANG) except EnvironmentError as err: if err.errno == EINTR: # retry handle_sigchld() elif err.errno != ECHILD: raise signal.signal(SIGCHLD, handle_sigchld) # Handle clients while not should_exit(): # Wait until a client arrives or we have to exit sock = None while not should_exit() and sock is None: try: sock, addr = listen_sock.accept() except EnvironmentError as err: if err.errno != EINTR: raise if sock is not None: # Fork a child to handle the client. # The client is handled in the child so that the manager # never receives SIGCHLD unless a worker crashes. if os.fork() == 0: # Leave the worker pool signal.signal(SIGHUP, SIG_DFL) listen_sock.close() # Read the socket using fdopen instead of socket.makefile() because the latter # seems to be very slow; note that we need to dup() the file descriptor because # otherwise writes also cause a seek that makes us miss data on the read side. infile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) outfile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) exit_code = 0 try: worker_main(infile, outfile) except SystemExit as exc: exit_code = exc.code finally: outfile.flush() sock.close() os._exit(compute_real_exit_code(exit_code)) else: sock.close()
def worker(listen_sock): # Redirect stdout to stderr os.dup2(2, 1) # Manager sends SIGHUP to request termination of workers in the pool def handle_sighup(*args): assert should_exit() signal(SIGHUP, handle_sighup) # Cleanup zombie children def handle_sigchld(*args): pid = status = None try: while (pid, status) != (0, 0): pid, status = os.waitpid(0, os.WNOHANG) except EnvironmentError as err: if err.errno == EINTR: # retry handle_sigchld() elif err.errno != ECHILD: raise signal(SIGCHLD, handle_sigchld) # Handle clients while not should_exit(): # Wait until a client arrives or we have to exit sock = None while not should_exit() and sock is None: try: sock, addr = listen_sock.accept() except EnvironmentError as err: if err.errno != EINTR: raise if sock is not None: # Fork a child to handle the client. # The client is handled in the child so that the manager # never receives SIGCHLD unless a worker crashes. if os.fork() == 0: # Leave the worker pool signal(SIGHUP, SIG_DFL) listen_sock.close() # Handle the client then exit sockfile = sock.makefile() exit_code = 0 try: worker_main(sockfile, sockfile) except SystemExit as exc: exit_code = exc.code finally: sockfile.close() sock.close() os._exit(compute_real_exit_code(exit_code)) else: sock.close()
def worker(sock): """ Called by a worker process after the fork(). """ # Redirect stdout to stderr os.dup2(2, 1) sys.stdout = sys.stderr # The sys.stdout object is different from file descriptor 1 signal.signal(SIGHUP, SIG_DFL) signal.signal(SIGCHLD, SIG_DFL) signal.signal(SIGTERM, SIG_DFL) # Blocks until the socket is closed by draining the input stream # until it raises an exception or returns EOF. def waitSocketClose(sock): try: while True: # Empty string is returned upon EOF (and only then). if sock.recv(4096) == '': return except: pass # Read the socket using fdopen instead of socket.makefile() because the latter # seems to be very slow; note that we need to dup() the file descriptor because # otherwise writes also cause a seek that makes us miss data on the read side. infile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) outfile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) exit_code = 0 try: # Acknowledge that the fork was successful write_int(os.getpid(), outfile) outfile.flush() worker_main(infile, outfile) except SystemExit as exc: exit_code = exc.code finally: outfile.flush() # The Scala side will close the socket upon task completion. waitSocketClose(sock) os._exit(compute_real_exit_code(exit_code))
def worker(sock, authenticated): """ Called by a worker process after the fork(). """ signal.signal(SIGHUP, SIG_DFL) signal.signal(SIGCHLD, SIG_DFL) signal.signal(SIGTERM, SIG_DFL) # restore the handler for SIGINT, # it's useful for debugging (show the stacktrace before exit) signal.signal(SIGINT, signal.default_int_handler) # Read the socket using fdopen instead of socket.makefile() because the latter # seems to be very slow; note that we need to dup() the file descriptor because # otherwise writes also cause a seek that makes us miss data on the read side. buffer_size = int(os.environ.get("SPARK_BUFFER_SIZE", 65536)) infile = os.fdopen(os.dup(sock.fileno()), "rb", buffer_size) outfile = os.fdopen(os.dup(sock.fileno()), "wb", buffer_size) if not authenticated: client_secret = UTF8Deserializer().loads(infile) if os.environ["PYTHON_WORKER_FACTORY_SECRET"] == client_secret: write_with_length("ok".encode("utf-8"), outfile) outfile.flush() else: write_with_length("err".encode("utf-8"), outfile) outfile.flush() sock.close() return 1 exit_code = 0 try: worker_main(infile, outfile) except SystemExit as exc: exit_code = compute_real_exit_code(exc.code) finally: try: outfile.flush() except Exception: pass return exit_code
def worker(sock, authenticated): """ Called by a worker process after the fork(). """ signal.signal(SIGHUP, SIG_DFL) signal.signal(SIGCHLD, SIG_DFL) signal.signal(SIGTERM, SIG_DFL) # restore the handler for SIGINT, # it's useful for debugging (show the stacktrace before exit) signal.signal(SIGINT, signal.default_int_handler) # Read the socket using fdopen instead of socket.makefile() because the latter # seems to be very slow; note that we need to dup() the file descriptor because # otherwise writes also cause a seek that makes us miss data on the read side. infile = os.fdopen(os.dup(sock.fileno()), "rb", 65536) outfile = os.fdopen(os.dup(sock.fileno()), "wb", 65536) if not authenticated: client_secret = UTF8Deserializer().loads(infile) if os.environ["PYTHON_WORKER_FACTORY_SECRET"] == client_secret: write_with_length("ok".encode("utf-8"), outfile) outfile.flush() else: write_with_length("err".encode("utf-8"), outfile) outfile.flush() sock.close() return 1 exit_code = 0 try: worker_main(infile, outfile) except SystemExit as exc: exit_code = compute_real_exit_code(exc.code) finally: try: outfile.flush() except Exception: pass return exit_code
def worker(sock): """ Called by a worker process after the fork(). """ signal.signal(SIGHUP, SIG_DFL) signal.signal(SIGCHLD, SIG_DFL) signal.signal(SIGTERM, SIG_DFL) # Read the socket using fdopen instead of socket.makefile() because the latter # seems to be very slow; note that we need to dup() the file descriptor because # otherwise writes also cause a seek that makes us miss data on the read side. infile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) outfile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) exit_code = 0 try: worker_main(infile, outfile) except SystemExit as exc: exit_code = compute_real_exit_code(exc.code) finally: outfile.flush() if exit_code: os._exit(exit_code)
print( "DEBUG: Pooled memory, pool size: {} MiB, max size: {} MiB".format( pool_size / 1024.0 / 1024, ('unlimited' if pool_max_size == max_size else pool_max_size / 1024.0 / 1024))) base_t = rmm.mr.ManagedMemoryResource if uvm_enabled else rmm.mr.CudaMemoryResource rmm.mr.set_current_device_resource( rmm.mr.PoolMemoryResource(base_t(), pool_size, pool_max_size)) elif uvm_enabled: from cudf import rmm rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) else: # Do nothing, whether to use RMM (default mode) or not depends on UDF definition. pass if __name__ == '__main__': # GPU context setup initialize_gpu_mem() # Code below is all copied from Pyspark/worker.py java_port = int(os.environ["PYTHON_WORKER_FACTORY_PORT"]) auth_secret = os.environ["PYTHON_WORKER_FACTORY_SECRET"] (sock_file, sock) = local_connect_and_auth(java_port, auth_secret) # Use the `sock_file` as both input and output will cause EOFException in JVM side, # So open a new file object on the same socket as output, similar behavior # with that in `pyspark/daemon.py`. buffer_size = int(os.environ.get("SPARK_BUFFER_SIZE", 65536)) outfile = os.fdopen(os.dup(sock.fileno()), "wb", buffer_size) worker_main(sock_file, outfile)
def worker(listen_sock): # Redirect stdout to stderr os.dup2(2, 1) sys.stdout = sys.stderr # The sys.stdout object is different from file descriptor 1 # Manager sends SIGHUP to request termination of workers in the pool def handle_sighup(*args): assert should_exit() signal.signal(SIGHUP, handle_sighup) # Cleanup zombie children def handle_sigchld(*args): pid = status = None try: while (pid, status) != (0, 0): pid, status = os.waitpid(0, os.WNOHANG) except EnvironmentError as err: if err.errno == EINTR: # retry handle_sigchld() elif err.errno != ECHILD: raise signal.signal(SIGCHLD, handle_sigchld) # Blocks until the socket is closed by draining the input stream # until it raises an exception or returns EOF. def waitSocketClose(sock): try: while True: # Empty string is returned upon EOF (and only then). if sock.recv(4096) == '': return except: pass # Handle clients while not should_exit(): # Wait until a client arrives or we have to exit sock = None while not should_exit() and sock is None: try: sock, addr = listen_sock.accept() except EnvironmentError as err: if err.errno != EINTR: raise if sock is not None: # Fork a child to handle the client. # The client is handled in the child so that the manager # never receives SIGCHLD unless a worker crashes. if os.fork() == 0: # Leave the worker pool signal.signal(SIGHUP, SIG_DFL) listen_sock.close() # Read the socket using fdopen instead of socket.makefile() because the latter # seems to be very slow; note that we need to dup() the file descriptor because # otherwise writes also cause a seek that makes us miss data on the read side. infile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) outfile = os.fdopen(os.dup(sock.fileno()), "a+", 65536) exit_code = 0 try: worker_main(infile, outfile) except SystemExit as exc: exit_code = exc.code finally: outfile.flush() # The Scala side will close the socket upon task completion. waitSocketClose(sock) os._exit(compute_real_exit_code(exit_code)) else: sock.close()