def test_grpc_socket_workspace(): first_server_process = GrpcServerProcess() with first_server_process.create_ephemeral_client() as first_server: second_server_process = GrpcServerProcess() with second_server_process.create_ephemeral_client() as second_server: first_socket = first_server.socket second_socket = second_server.socket workspace_yaml = """ load_from: - grpc_server: host: localhost socket: {socket_one} - grpc_server: socket: {socket_two} location_name: 'local_port_default_host' """.format(socket_one=first_socket, socket_two=second_socket) origins = location_origins_from_config( yaml.safe_load(workspace_yaml), # fake out as if it were loaded by a yaml file in this directory file_relative_path(__file__, "not_a_real.yaml"), ) with ExitStack() as stack: repository_location_handles = { name: stack.enter_context(origin.create_handle()) for name, origin in origins.items() } assert len(repository_location_handles) == 2 default_location_name = "grpc:localhost:{socket}".format( socket=first_socket) assert repository_location_handles.get(default_location_name) local_port = repository_location_handles.get( default_location_name) assert local_port.socket == first_socket assert local_port.host == "localhost" assert local_port.port is None assert repository_location_handles.get( "local_port_default_host") local_port_default_host = repository_location_handles.get( "local_port_default_host") assert local_port_default_host.socket == second_socket assert local_port_default_host.host == "localhost" assert local_port_default_host.port is None assert all( map(lambda x: x.location_name, repository_location_handles.values())) second_server_process.wait() first_server_process.wait()
def test_grpc_socket_workspace(): first_server_process = GrpcServerProcess() with first_server_process.create_ephemeral_client() as first_server: second_server_process = GrpcServerProcess() with second_server_process.create_ephemeral_client() as second_server: first_socket = first_server.socket second_socket = second_server.socket workspace_yaml = """ load_from: - grpc_server: host: localhost socket: {socket_one} - grpc_server: socket: {socket_two} location_name: 'local_port_default_host' """.format(socket_one=first_socket, socket_two=second_socket) workspace = load_workspace_from_config( yaml.safe_load(workspace_yaml), # fake out as if it were loaded by a yaml file in this directory file_relative_path(__file__, "not_a_real.yaml"), ) assert isinstance(workspace, Workspace) assert len(workspace.repository_location_handles) == 2 default_location_name = "grpc:localhost:{socket}".format( socket=first_socket) assert workspace.has_repository_location_handle( default_location_name) local_port = workspace.get_repository_location_handle( default_location_name) assert local_port.socket == first_socket assert local_port.host == "localhost" assert local_port.port is None assert workspace.has_repository_location_handle( "local_port_default_host") local_port_default_host = workspace.get_repository_location_handle( "local_port_default_host") assert local_port_default_host.socket == second_socket assert local_port_default_host.host == "localhost" assert local_port_default_host.port is None assert all( map(lambda x: x.location_name, workspace.repository_location_handles)) second_server_process.wait() first_server_process.wait()
def test_server_down(): with grpc_instance() as instance: repo_yaml = file_relative_path(__file__, "repo.yaml") recon_repo = ReconstructableRepository.from_legacy_repository_yaml( repo_yaml) loadable_target_origin = recon_repo.get_origin().loadable_target_origin server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=4, force_port=True) with server_process.create_ephemeral_client() as api_client: repository_location = GrpcServerRepositoryLocation( RepositoryLocationHandle.create_grpc_server_location( location_name="test", port=api_client.port, socket=api_client.socket, host=api_client.host, )) external_pipeline = repository_location.get_repository( "nope").get_full_external_pipeline("sleepy_pipeline") pipeline_run = instance.create_run_for_pipeline( pipeline_def=sleepy_pipeline, run_config=None) launcher = instance.run_launcher launcher.launch_run(instance, pipeline_run, external_pipeline) poll_for_step_start(instance, pipeline_run.run_id) assert launcher.can_terminate(pipeline_run.run_id) original_run_tags = instance.get_run_by_id( pipeline_run.run_id).tags[GRPC_INFO_TAG] # Replace run tags with an invalid port instance.add_run_tags( pipeline_run.run_id, { GRPC_INFO_TAG: seven.json.dumps( merge_dicts({"host": "localhost"}, {"port": find_free_port()})) }, ) assert not launcher.can_terminate(pipeline_run.run_id) instance.add_run_tags( pipeline_run.run_id, { GRPC_INFO_TAG: original_run_tags, }, ) assert launcher.terminate(pipeline_run.run_id) server_process.wait()
def test_server_down(): with instance_for_test() as instance: loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, attribute="nope", python_file=file_relative_path(__file__, "test_default_run_launcher.py"), ) server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=4, force_port=True ) with server_process.create_ephemeral_client() as api_client: repository_location = GrpcServerRepositoryLocation( RepositoryLocationHandle.create_from_repository_location_origin( GrpcServerRepositoryLocationOrigin( location_name="test", port=api_client.port, socket=api_client.socket, host=api_client.host, ) ) ) external_pipeline = repository_location.get_repository( "nope" ).get_full_external_pipeline("sleepy_pipeline") pipeline_run = instance.create_run_for_pipeline( pipeline_def=sleepy_pipeline, run_config=None ) instance.launch_run(pipeline_run.run_id, external_pipeline) poll_for_step_start(instance, pipeline_run.run_id) launcher = instance.run_launcher assert launcher.can_terminate(pipeline_run.run_id) original_run_tags = instance.get_run_by_id(pipeline_run.run_id).tags[GRPC_INFO_TAG] # Replace run tags with an invalid port instance.add_run_tags( pipeline_run.run_id, { GRPC_INFO_TAG: seven.json.dumps( merge_dicts({"host": "localhost"}, {"port": find_free_port()}) ) }, ) assert not launcher.can_terminate(pipeline_run.run_id) instance.add_run_tags( pipeline_run.run_id, {GRPC_INFO_TAG: original_run_tags,}, ) assert launcher.terminate(pipeline_run.run_id) server_process.wait()
def get_deployed_grpc_server_workspace(instance): loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, attribute="nope", python_file=file_relative_path(__file__, "test_default_run_launcher.py"), ) server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin) try: with server_process.create_ephemeral_client( ): # shuts down when leaves this context with WorkspaceProcessContext( instance, GrpcServerTarget( host="localhost", socket=server_process.socket, port=server_process.port, location_name="test", ), ) as workspace_process_context: yield workspace_process_context.create_request_context() finally: server_process.wait()
def create_process_bound_grpc_server_location(loadable_target_origin, location_name): from dagster.grpc.client import client_heartbeat_thread from dagster.grpc.server import GrpcServerProcess server = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=2, heartbeat=True) client = server.create_ephemeral_client() heartbeat_thread = threading.Thread(target=client_heartbeat_thread, args=(client, )) heartbeat_thread.daemon = True heartbeat_thread.start() list_repositories_response = sync_list_repositories_grpc(client) code_pointer_dict = list_repositories_response.repository_code_pointer_dict return ManagedGrpcPythonEnvRepositoryLocationHandle( executable_path=list_repositories_response.executable_path, location_name=location_name if location_name else _assign_python_env_location_name(code_pointer_dict), repository_code_pointer_dict=code_pointer_dict, client=client, grpc_server_process=server, )
def get_external_pipeline_from_grpc_server_repository(pipeline_name): loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, attribute="nope", python_file=file_relative_path(__file__, "test_default_run_launcher.py"), ) server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin) try: with server_process.create_ephemeral_client() as api_client: repository_location = GrpcServerRepositoryLocation( RepositoryLocationHandle. create_from_repository_location_origin( GrpcServerRepositoryLocationOrigin( location_name="test", port=api_client.port, socket=api_client.socket, host=api_client.host, ))) yield repository_location.get_repository( "nope").get_full_external_pipeline(pipeline_name) finally: server_process.wait()
def test_heartbeat(): loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, attribute="bar_repo", python_file=file_relative_path(__file__, "grpc_repo.py"), ) server = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=2, heartbeat=True, heartbeat_timeout=1, ) with server.create_ephemeral_client() as client: assert server.server_process.poll() is None # heartbeat keeps the server alive time.sleep(0.5) client.heartbeat() time.sleep(0.5) client.heartbeat() time.sleep(0.5) assert server.server_process.poll() is None start_time = time.time() while (time.time() - start_time) < 10: if server.server_process.poll() is not None: return time.sleep(0.1) raise Exception( "Timed out waiting for server to terminate after heartbeat stopped" )
def test_heartbeat(): loadable_target_origin = LoadableTargetOrigin( attribute="bar_repo", python_file=file_relative_path(__file__, "grpc_repo.py"), ) server = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=2, heartbeat=True, heartbeat_timeout=1, ) with server.create_ephemeral_client() as client: assert server.server_process.poll() is None # heartbeat keeps the server alive time.sleep(0.5) client.heartbeat() time.sleep(0.5) client.heartbeat() time.sleep(0.5) assert server.server_process.poll() is None # without the heartbeat, the server dies time.sleep(2) assert server.server_process.poll() is not None
def test_ssl_grpc_server_workspace(): server_process = GrpcServerProcess(force_port=True) try: with server_process.create_ephemeral_client() as client: assert client.heartbeat(echo="Hello") port = server_process.port ssl_yaml = f""" load_from: - grpc_server: host: localhost port: {port} ssl: true """ origins = location_origins_from_config( yaml.safe_load(ssl_yaml), # fake out as if it were loaded by a yaml file in this directory file_relative_path(__file__, "not_a_real.yaml"), ) origin = list(origins.values())[0] assert origin.use_ssl # Actually connecting to the server will fail since it's expecting SSL # and we didn't set up the server with SSL try: with origin.create_location(): assert False except DagsterUserCodeUnreachableError: pass finally: server_process.wait()
def test_list_command_grpc_socket(): with instance_for_test() as instance: runner = CliRunner() server_process = GrpcServerProcess( loadable_target_origin=LoadableTargetOrigin( executable_path=sys.executable, python_file=file_relative_path(__file__, "test_cli_commands.py"), attribute="bar", ), ) with server_process.create_ephemeral_client() as api_client: execute_list_command( {"grpc_socket": api_client.socket}, no_print, instance, ) execute_list_command( {"grpc_socket": api_client.socket, "grpc_host": api_client.host}, no_print, instance, ) result = runner.invoke(pipeline_list_command, ["--grpc-socket", api_client.socket]) assert_correct_bar_repository_output(result) result = runner.invoke( pipeline_list_command, ["--grpc-socket", api_client.socket, "--grpc-host", api_client.host], ) assert_correct_bar_repository_output(result) server_process.wait()
def test_run_always_finishes(): # pylint: disable=redefined-outer-name with seven.TemporaryDirectory() as temp_dir: instance = DagsterInstance.local_temp( temp_dir, overrides={ "run_launcher": { "module": "dagster.core.launcher.grpc_run_launcher", "class": "GrpcRunLauncher", } }, ) pipeline_run = instance.create_run_for_pipeline( pipeline_def=slow_pipeline, run_config=None) run_id = pipeline_run.run_id loadable_target_origin = LoadableTargetOrigin( attribute="nope", python_file=file_relative_path(__file__, "test_default_run_launcher.py"), ) server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=4) with server_process.create_ephemeral_client() as api_client: repository_location = GrpcServerRepositoryLocation( RepositoryLocationHandle.create_grpc_server_location( location_name="test", port=api_client.port, socket=api_client.socket, host=api_client.host, )) external_pipeline = repository_location.get_repository( "nope").get_full_external_pipeline("slow_pipeline") assert instance.get_run_by_id( run_id).status == PipelineRunStatus.NOT_STARTED launcher = instance.run_launcher launcher.launch_run(instance=instance, run=pipeline_run, external_pipeline=external_pipeline) # Server process now receives shutdown event, run has not finished yet pipeline_run = instance.get_run_by_id(run_id) assert not pipeline_run.is_finished assert server_process.server_process.poll() is None # Server should wait until run finishes, then shutdown pipeline_run = poll_for_finished_run(instance, run_id) assert pipeline_run.status == PipelineRunStatus.SUCCESS start_time = time.time() while server_process.server_process.poll() is None: time.sleep(0.05) # Verify server process cleans up eventually assert time.time() - start_time < 5 server_process.wait()
def test_list_command_deployed_grpc(): runner = CliRunner() with instance_for_test() as instance: server_process = GrpcServerProcess( loadable_target_origin=LoadableTargetOrigin( executable_path=sys.executable, python_file=file_relative_path(__file__, "test_cli_commands.py"), attribute="bar", ), force_port=True, ) with server_process.create_ephemeral_client() as api_client: result = runner.invoke(pipeline_list_command, ["--grpc-port", api_client.port]) assert_correct_bar_repository_output(result) result = runner.invoke( pipeline_list_command, [ "--grpc-port", api_client.port, "--grpc-host", api_client.host ], ) assert_correct_bar_repository_output(result) result = runner.invoke(pipeline_list_command, ["--grpc-port", api_client.port]) assert_correct_bar_repository_output(result) result = runner.invoke( pipeline_list_command, [ "--grpc-port", api_client.port, "--grpc-socket", "foonamedsocket" ], ) assert result.exit_code != 0 execute_list_command( {"grpc_port": api_client.port}, no_print, instance, ) # Can't supply both port and socket with pytest.raises(UsageError): execute_list_command( { "grpc_port": api_client.port, "grpc_socket": "foonamedsocket" }, no_print, instance, ) server_process.wait()
def test_cancel_run(): with instance_for_test() as instance: loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, python_file=__file__, working_directory=None, ) server_process = GrpcServerProcess(loadable_target_origin, max_workers=10) with server_process.create_ephemeral_client() as api_client: streaming_results = [] pipeline_run = instance.create_run_for_pipeline( streaming_pipeline, run_config={"solids": {"streamer": {"config": {"length": 20}}}}, ) execute_run_args = ExecuteRunArgs( pipeline_origin=PipelineGrpcServerOrigin( pipeline_name="streaming_pipeline", repository_origin=RepositoryGrpcServerOrigin( host="localhost", socket=api_client.socket, port=api_client.port, repository_name="test_repository", ), ), pipeline_run_id=pipeline_run.run_id, instance_ref=instance.get_ref(), ) stream_events_result_thread = threading.Thread( target=_stream_events_target, args=[streaming_results, api_client, execute_run_args] ) stream_events_result_thread.daemon = True stream_events_result_thread.start() poll_for_step_start(instance, pipeline_run.run_id) res = api_client.cancel_execution( cancel_execution_request=CancelExecutionRequest(run_id=pipeline_run.run_id) ) assert res.success is True poll_for_run(instance, pipeline_run.run_id) logs = instance.all_logs(pipeline_run.run_id) assert ( len( [ ev for ev in logs if ev.dagster_event.event_type_value == "STEP_MATERIALIZATION" ] ) < 20 ) # soft termination assert [ev for ev in logs if ev.dagster_event.event_type_value == "STEP_FAILURE"] server_process.wait()
def test_run_always_finishes(): # pylint: disable=redefined-outer-name with instance_for_test() as instance: loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, attribute="nope", python_file=file_relative_path(__file__, "test_default_run_launcher.py"), ) server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=4 ) with server_process.create_ephemeral_client(): # Shuts down when leaves context with WorkspaceProcessContext( instance, GrpcServerTarget( host="localhost", socket=server_process.socket, port=server_process.port, location_name="test", ), ) as workspace_process_context: workspace = workspace_process_context.create_request_context() external_pipeline = ( workspace.get_repository_location("test") .get_repository("nope") .get_full_external_pipeline("slow_pipeline") ) pipeline_run = instance.create_run_for_pipeline( pipeline_def=slow_pipeline, run_config=None, external_pipeline_origin=external_pipeline.get_external_origin(), pipeline_code_origin=external_pipeline.get_python_origin(), ) run_id = pipeline_run.run_id assert instance.get_run_by_id(run_id).status == PipelineRunStatus.NOT_STARTED instance.launch_run(run_id=run_id, workspace=workspace) # Server process now receives shutdown event, run has not finished yet pipeline_run = instance.get_run_by_id(run_id) assert not pipeline_run.is_finished assert server_process.server_process.poll() is None # Server should wait until run finishes, then shutdown pipeline_run = poll_for_finished_run(instance, run_id) assert pipeline_run.status == PipelineRunStatus.SUCCESS start_time = time.time() while server_process.server_process.poll() is None: time.sleep(0.05) # Verify server process cleans up eventually assert time.time() - start_time < 5 server_process.wait()
def test_execute_run_api_grpc_python_handle(): with instance_for_test() as instance: with get_foo_pipeline_handle() as pipeline_handle: pipeline_run = instance.create_run( pipeline_name="foo", run_id=None, run_config={}, mode="default", solids_to_execute=None, step_keys_to_execute=None, status=None, tags=None, root_run_id=None, parent_run_id=None, pipeline_snapshot=None, execution_plan_snapshot=None, parent_pipeline_snapshot=None, ) loadable_target_origin = (pipeline_handle.get_external_origin( ).external_repository_origin.repository_location_origin. loadable_target_origin) server_process = GrpcServerProcess(loadable_target_origin, max_workers=2) with server_process.create_ephemeral_client() as api_client: events = [ event for event in sync_execute_run_grpc( api_client=api_client, instance_ref=instance.get_ref(), pipeline_origin=pipeline_handle.get_external_origin(), pipeline_run=pipeline_run, ) ] assert len(events) == 17 assert [event.event_type_value for event in events] == [ "ENGINE_EVENT", "ENGINE_EVENT", "PIPELINE_START", "ENGINE_EVENT", "STEP_START", "STEP_OUTPUT", "OBJECT_STORE_OPERATION", "STEP_SUCCESS", "STEP_START", "OBJECT_STORE_OPERATION", "STEP_INPUT", "STEP_OUTPUT", "OBJECT_STORE_OPERATION", "STEP_SUCCESS", "ENGINE_EVENT", "PIPELINE_SUCCESS", "ENGINE_EVENT", ] server_process.wait()
def test_health_check_success(): loadable_target_origin = LoadableTargetOrigin( attribute="bar_repo", python_file=file_relative_path(__file__, "grpc_repo.py"), ) server = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=2, heartbeat=True, heartbeat_timeout=1, ) with server.create_ephemeral_client() as client: assert client.health_check_query() == "SERVING"
def grpc_repo_location(): loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, python_file=__file__, attribute="the_repo" ) server_process = GrpcServerProcess(loadable_target_origin=loadable_target_origin) try: with server_process.create_ephemeral_client() as api_client: yield RepositoryLocation.from_handle( RepositoryLocationHandle.create_grpc_server_location( port=api_client.port, socket=api_client.socket, host=api_client.host, ) ) finally: server_process.wait()
def grpc_schedule_origin(schedule_name): loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, python_file=__file__, attribute="the_repo" ) server_process = GrpcServerProcess(loadable_target_origin=loadable_target_origin) with server_process.create_ephemeral_client() as api_client: repo_origin = RepositoryGrpcServerOrigin( host=api_client.host, port=api_client.port, socket=api_client.socket, repository_name="the_repo", ) yield repo_origin.get_schedule_origin(schedule_name) server_process.wait()
def test_run_always_finishes(): # pylint: disable=redefined-outer-name with instance_for_test() as instance: pipeline_run = instance.create_run_for_pipeline( pipeline_def=slow_pipeline, run_config=None) run_id = pipeline_run.run_id loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, attribute="nope", python_file=file_relative_path(__file__, "test_default_run_launcher.py"), ) server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=4) with server_process.create_ephemeral_client() as api_client: repository_location = GrpcServerRepositoryLocation( RepositoryLocationHandle. create_from_repository_location_origin( GrpcServerRepositoryLocationOrigin( location_name="test", port=api_client.port, socket=api_client.socket, host=api_client.host, ))) external_pipeline = repository_location.get_repository( "nope").get_full_external_pipeline("slow_pipeline") assert instance.get_run_by_id( run_id).status == PipelineRunStatus.NOT_STARTED instance.launch_run(run_id=pipeline_run.run_id, external_pipeline=external_pipeline) # Server process now receives shutdown event, run has not finished yet pipeline_run = instance.get_run_by_id(run_id) assert not pipeline_run.is_finished assert server_process.server_process.poll() is None # Server should wait until run finishes, then shutdown pipeline_run = poll_for_finished_run(instance, run_id) assert pipeline_run.status == PipelineRunStatus.SUCCESS start_time = time.time() while server_process.server_process.poll() is None: time.sleep(0.05) # Verify server process cleans up eventually assert time.time() - start_time < 5 server_process.wait()
def grpc_server_bar_kwargs(pipeline_name=None): server_process = GrpcServerProcess( loadable_target_origin=LoadableTargetOrigin( python_file=file_relative_path(__file__, "test_cli_commands.py"), attribute="bar"), ) with server_process.create_ephemeral_client() as client: args = {"grpc_host": client.host} if pipeline_name: args["pipeline"] = "foo" if client.port: args["grpc_port"] = client.port if client.socket: args["grpc_socket"] = client.socket yield args server_process.wait()
def test_heartbeat(): loadable_target_origin = LoadableTargetOrigin( attribute='bar_repo', python_file=file_relative_path(__file__, 'grpc_repo.py'), ) server = GrpcServerProcess( loadable_target_origin=loadable_target_origin, max_workers=2, heartbeat=True, heartbeat_timeout=1, ) with server.create_ephemeral_client() as client: client.heartbeat() assert server.server_process.poll() is None time.sleep(2) assert server.server_process.poll() is not None
def test_process_killed_after_client_finished(): server_process = GrpcServerProcess() with server_process.create_ephemeral_client() as client: socket = client.socket assert socket and os.path.exists(socket) start_time = time.time() while server_process.server_process.poll() is None: time.sleep(0.05) # Verify server process cleans up eventually assert time.time() - start_time < 1.5 # verify socket is cleaned up assert not os.path.exists(socket)
def test_run_always_finishes(temp_instance): # pylint: disable=redefined-outer-name instance = temp_instance pipeline_run = instance.create_run_for_pipeline(pipeline_def=slow_pipeline, run_config=None) run_id = pipeline_run.run_id recon_repo = ReconstructableRepository.for_file(__file__, 'nope') loadable_target_origin = LoadableTargetOrigin.from_python_origin(recon_repo.get_origin()) server_process = GrpcServerProcess(loadable_target_origin=loadable_target_origin, max_workers=4) with server_process.create_ephemeral_client() as api_client: repository_location = GrpcServerRepositoryLocation( RepositoryLocationHandle.create_grpc_server_location( location_name='test', port=api_client.port, socket=api_client.socket, host=api_client.host, ) ) external_pipeline = repository_location.get_repository('nope').get_full_external_pipeline( 'slow_pipeline' ) assert instance.get_run_by_id(run_id).status == PipelineRunStatus.NOT_STARTED launcher = instance.run_launcher launcher.launch_run( instance=instance, run=pipeline_run, external_pipeline=external_pipeline ) # Server process now receives shutdown event, run has not finished yet pipeline_run = instance.get_run_by_id(run_id) assert not pipeline_run.is_finished assert server_process.server_process.poll() is None # Server should wait until run finishes, then shutdown pipeline_run = poll_for_run(instance, run_id) assert pipeline_run.status == PipelineRunStatus.SUCCESS start_time = time.time() while server_process.server_process.poll() is None: time.sleep(0.05) # Verify server process cleans up eventually assert time.time() - start_time < 5
def _mgr_fn(instance, read_only): server_process = GrpcServerProcess( loadable_target_origin=get_main_loadable_target_origin()) try: with server_process.create_ephemeral_client() as api_client: with WorkspaceProcessContext( instance, GrpcServerTarget( port=api_client.port, socket=api_client.socket, host=api_client.host, location_name="test", ), version="", read_only=read_only, ) as workspace: yield workspace finally: server_process.wait()
def grpc_server_bar_cli_args(pipeline_name=None): server_process = GrpcServerProcess( loadable_target_origin=LoadableTargetOrigin( python_file=file_relative_path(__file__, "test_cli_commands.py"), attribute="bar"), ) with server_process.create_ephemeral_client() as client: args = ["--grpc-host", client.host] if client.port: args.append("--grpc-port") args.append(client.port) if client.socket: args.append("--grpc-socket") args.append(client.socket) if pipeline_name: args.append("--pipeline") args.append(pipeline_name) yield args server_process.wait()
def workspace_process_context_fixture(instance): loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, python_file=file_relative_path(__file__, "test_custom_repository_data.py"), ) server_process = GrpcServerProcess(loadable_target_origin=loadable_target_origin) try: with server_process.create_ephemeral_client(): # shuts down when leaves this context with WorkspaceProcessContext( instance, GrpcServerTarget( host="localhost", socket=server_process.socket, port=server_process.port, location_name="test", ), ) as workspace_process_context: yield workspace_process_context finally: server_process.wait()
def _mgr_fn(recon_repo): check.inst_param(recon_repo, "recon_repo", ReconstructableRepository) loadable_target_origin = recon_repo.get_python_origin( ).loadable_target_origin server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin) try: with server_process.create_ephemeral_client() as api_client: with Workspace( GrpcServerTarget( port=api_client.port, socket=api_client.socket, host=api_client.host, location_name="test", )) as workspace: yield workspace finally: server_process.wait()
def get_external_pipeline_from_grpc_server_repository(pipeline_name): repo_yaml = file_relative_path(__file__, "repo.yaml") recon_repo = ReconstructableRepository.from_legacy_repository_yaml( repo_yaml) loadable_target_origin = recon_repo.get_origin().loadable_target_origin server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin) try: with server_process.create_ephemeral_client() as api_client: repository_location = GrpcServerRepositoryLocation( RepositoryLocationHandle.create_grpc_server_location( location_name="test", port=api_client.port, socket=api_client.socket, host=api_client.host, )) yield repository_location.get_repository( "nope").get_full_external_pipeline(pipeline_name) finally: server_process.wait()
class ManagedGrpcPythonEnvRepositoryLocationHandle(RepositoryLocationHandle): """ A Python environment for which Dagster is managing a gRPC server. """ def __init__(self, origin): from dagster.grpc.client import client_heartbeat_thread from dagster.grpc.server import GrpcServerProcess self.grpc_server_process = None self.client = None self.heartbeat_shutdown_event = None self.heartbeat_thread = None self._origin = check.inst_param( origin, "origin", ManagedGrpcPythonEnvRepositoryLocationOrigin) loadable_target_origin = origin.loadable_target_origin self._external_repositories_data = None try: self.grpc_server_process = GrpcServerProcess( loadable_target_origin=loadable_target_origin, heartbeat=True, ) self.client = self.grpc_server_process.create_ephemeral_client() self.heartbeat_shutdown_event = threading.Event() self.heartbeat_thread = threading.Thread( target=client_heartbeat_thread, args=( self.client, self.heartbeat_shutdown_event, ), name="grpc-client-heartbeat", ) self.heartbeat_thread.daemon = True self.heartbeat_thread.start() list_repositories_response = sync_list_repositories_grpc( self.client) self.repository_code_pointer_dict = ( list_repositories_response.repository_code_pointer_dict) self.container_image = self.client.get_current_image( ).current_image self._external_repositories_data = sync_get_streaming_external_repositories_data_grpc( self.client, self, ) except: self.cleanup() raise def create_external_repositories(self): from dagster.core.host_representation.external import ExternalRepository return { repo_name: ExternalRepository( repo_data, RepositoryHandle( repository_name=repo_name, repository_location_handle=self, ), ) for repo_name, repo_data in self._external_repositories_data.items() } def get_repository_python_origin(self, repository_name): return _get_repository_python_origin( self.executable_path, self.repository_code_pointer_dict, repository_name, self.container_image, ) @property def origin(self): return self._origin @property def executable_path(self): return self.loadable_target_origin.executable_path @property def location_name(self): return self.origin.location_name @property def loadable_target_origin(self): return self.origin.loadable_target_origin @property def repository_names(self): return set(self.repository_code_pointer_dict.keys()) @property def host(self): return "localhost" @property def port(self): return self.grpc_server_process.port @property def socket(self): return self.grpc_server_process.socket @property def use_ssl(self): return False def cleanup(self): if self.heartbeat_shutdown_event: self.heartbeat_shutdown_event.set() self.heartbeat_shutdown_event = None if self.heartbeat_thread: self.heartbeat_thread.join() self.heartbeat_thread = None if self.client: self.client.cleanup_server() self.client = None @property def is_cleaned_up(self): return not self.client def create_location(self): from dagster.core.host_representation.repository_location import ( GrpcServerRepositoryLocation, ) return GrpcServerRepositoryLocation(self) def get_display_metadata(self): return merge_dicts( self.origin.get_display_metadata(), ({ "image": self.container_image } if self.container_image else {}), )