def _mgr_fn(recon_repo):
"""Goes out of process but same process as host process"""
check.inst_param(recon_repo, "recon_repo",
ReconstructableRepository)
empty_repo = ReconstructableRepository.from_legacy_repository_yaml(
file_relative_path(__file__, "empty_repo.yaml"))
with Workspace([
RepositoryLocationHandle.create_python_env_location(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
python_file=file_relative_path(
__file__, "setup.py"),
attribute="test_repo",
),
location_name="test",
),
RepositoryLocationHandle.create_in_process_location(
empty_repo.pointer),
RepositoryLocationHandle.create_python_env_location(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
python_file=file_relative_path(
__file__, "setup.py"),
attribute="empty_repo",
),
location_name="empty_repo",
),
]) as workspace:
yield workspace
def _mgr_fn(recon_repo):
"""Goes out of process but same process as host process"""
check.inst_param(recon_repo, "recon_repo", ReconstructableRepository)
with Workspace(
[
ManagedGrpcPythonEnvRepositoryLocationOrigin(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
python_file=file_relative_path(__file__, "setup.py"),
attribute="test_repo",
),
location_name="test",
),
ManagedGrpcPythonEnvRepositoryLocationOrigin(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
python_file=file_relative_path(__file__, "setup.py"),
attribute="empty_repo",
),
location_name="empty_repo",
),
]
) as workspace:
yield workspace
def test_raise_on_error(kubeconfig_file):
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.MagicMock()
k8s_run_launcher = K8sRunLauncher(
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
job_image="fake_job_image",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
fail_pod_on_run_failure=True,
)
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
repo_def = recon_repo.get_definition()
loadable_target_origin = LoadableTargetOrigin(python_file=__file__)
with instance_for_test() as instance:
with in_process_test_workspace(instance, loadable_target_origin) as workspace:
location = workspace.get_repository_location(workspace.repository_location_names[0])
repo_handle = RepositoryHandle(
repository_name=repo_def.name,
repository_location=location,
)
fake_external_pipeline = external_pipeline_from_recon_pipeline(
recon_pipeline,
solid_selection=None,
repository_handle=repo_handle,
)
# Launch the run in a fake Dagster instance.
pipeline_name = "demo_pipeline"
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
external_pipeline_origin=fake_external_pipeline.get_external_origin(),
pipeline_code_origin=fake_external_pipeline.get_python_origin(),
)
k8s_run_launcher.register_instance(instance)
k8s_run_launcher.launch_run(LaunchRunContext(run, workspace))
mock_method_calls = mock_k8s_client_batch_api.method_calls
assert len(mock_method_calls) > 0
method_name, _args, kwargs = mock_method_calls[0]
assert method_name == "create_namespaced_job"
container = kwargs["body"].spec.template.spec.containers[0]
args = container.args
assert (
args
== ExecuteRunArgs(
pipeline_origin=run.pipeline_code_origin,
pipeline_run_id=run.run_id,
instance_ref=instance.get_ref(),
set_exit_code_on_failure=True,
).get_command_args()
)
def test_process_server_registry():
origin = ManagedGrpcPythonEnvRepositoryLocationOrigin(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
attribute="repo",
python_file=file_relative_path(__file__,
"test_grpc_server_registry.py"),
), )
with ProcessGrpcServerRegistry(reload_interval=5,
heartbeat_ttl=10,
startup_timeout=5) as registry:
endpoint_one = registry.get_grpc_endpoint(origin)
endpoint_two = registry.get_grpc_endpoint(origin)
assert endpoint_two == endpoint_one
assert _can_connect(origin, endpoint_one)
assert _can_connect(origin, endpoint_two)
start_time = time.time()
while True:
# Registry should return a new server endpoint after 5 seconds
endpoint_three = registry.get_grpc_endpoint(origin)
if endpoint_three.server_id != endpoint_one.server_id:
break
if time.time() - start_time > 15:
raise Exception("Server ID never changed")
time.sleep(1)
assert _can_connect(origin, endpoint_three)
start_time = time.time()
while True:
# Server at endpoint_one should eventually die due to heartbeat failure
if not _can_connect(origin, endpoint_one):
break
if time.time() - start_time > 30:
raise Exception(
"Old Server never died after process manager released it")
time.sleep(1)
# Make one more fresh process, then leave the context so that it will be cleaned up
while True:
endpoint_four = registry.get_grpc_endpoint(origin)
if endpoint_four.server_id != endpoint_three.server_id:
assert _can_connect(origin, endpoint_four)
break
registry.wait_for_processes()
assert not _can_connect(origin, endpoint_three)
assert not _can_connect(origin, endpoint_four)
def grpc_command(
port=None,
socket=None,
host="localhost",
max_workers=None,
heartbeat=False,
heartbeat_timeout=30,
lazy_load_user_code=False,
ipc_output_file=None,
fixed_server_id=None,
override_system_timezone=None,
**kwargs,
):
if seven.IS_WINDOWS and port is None:
raise click.UsageError(
"You must pass a valid --port/-p on Windows: --socket/-s not supported."
)
if not (port or socket and not (port and socket)):
raise click.UsageError("You must pass one and only one of --port/-p or --socket/-s.")
loadable_target_origin = None
if any(
kwargs[key]
for key in [
"attribute",
"working_directory",
"module_name",
"package_name",
"python_file",
"empty_working_directory",
]
):
loadable_target_origin = LoadableTargetOrigin(
executable_path=sys.executable,
attribute=kwargs["attribute"],
working_directory=get_working_directory_from_kwargs(kwargs),
module_name=kwargs["module_name"],
python_file=kwargs["python_file"],
package_name=kwargs["package_name"],
)
with (
mock_system_timezone(override_system_timezone)
if override_system_timezone
else nullcontext()
):
server = DagsterGrpcServer(
port=port,
socket=socket,
host=host,
loadable_target_origin=loadable_target_origin,
max_workers=max_workers,
heartbeat=heartbeat,
heartbeat_timeout=heartbeat_timeout,
lazy_load_user_code=lazy_load_user_code,
ipc_output_file=ipc_output_file,
fixed_server_id=fixed_server_id,
)
server.serve()
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_error_repo_in_registry():
error_origin = ManagedGrpcPythonEnvRepositoryLocationOrigin(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
attribute="error_repo",
python_file=file_relative_path(__file__, "error_repo.py"),
),
)
with ProcessGrpcServerRegistry(reload_interval=5, heartbeat_ttl=10) as registry:
# Repository with a loading error does not raise an exception
endpoint = registry.get_grpc_endpoint(error_origin)
# But using that endpoint to load a location results in an error
with pytest.raises(DagsterUserCodeProcessError, match="object is not callable"):
with GrpcServerRepositoryLocation(
origin=error_origin,
server_id=endpoint.server_id,
port=endpoint.port,
socket=endpoint.socket,
host=endpoint.host,
watch_server=False,
):
pass
# that error is idempotent
with pytest.raises(DagsterUserCodeProcessError, match="object is not callable"):
with GrpcServerRepositoryLocation(
origin=error_origin,
server_id=endpoint.server_id,
port=endpoint.port,
socket=endpoint.socket,
host=endpoint.host,
watch_server=False,
):
pass
def test_list_command_grpc_socket():
runner = CliRunner()
with GrpcServerProcess(
loadable_target_origin=LoadableTargetOrigin(
python_file=file_relative_path(__file__, "test_cli_commands.py"), attribute="bar"
),
).create_ephemeral_client() as api_client:
execute_list_command(
{"grpc_socket": api_client.socket}, no_print, DagsterInstance.local_temp(),
)
execute_list_command(
{"grpc_socket": api_client.socket, "grpc_host": api_client.host},
no_print,
DagsterInstance.local_temp(),
)
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)
def location_handle_from_python_file(
python_file,
attribute,
working_directory,
user_process_api,
location_name=None,
executable_path=sys.executable,
):
check.str_param(python_file, "python_file")
check.opt_str_param(attribute, "attribute")
check.opt_str_param(working_directory, "working_directory")
check.inst_param(user_process_api, "user_process_api", UserProcessApi)
check.opt_str_param(location_name, "location_name")
loadable_target_origin = LoadableTargetOrigin(
executable_path=executable_path,
python_file=python_file,
module_name=None,
working_directory=working_directory,
attribute=attribute,
)
return RepositoryLocationHandle.create_python_env_location(
loadable_target_origin=loadable_target_origin,
location_name=location_name,
user_process_api=user_process_api,
)
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 loadable_target_origin():
return LoadableTargetOrigin(
executable_path=sys.executable,
module_name="dagster_tests.scheduler_tests.test_scheduler_run",
working_directory=os.getcwd(),
attribute="the_repo",
)
def get_loadable_target_origin(self, executable_path):
return LoadableTargetOrigin(
executable_path=executable_path,
python_file=self.python_file,
attribute=self.fn_name,
working_directory=self.working_directory,
)
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 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_bad_load():
with _default_instance() as instance:
instance = DagsterInstance.get()
working_directory = os.path.dirname(__file__)
loadable_target_origin = LoadableTargetOrigin(
executable_path=sys.executable,
python_file=__file__,
attribute="doesnt_exist",
working_directory=working_directory,
)
repo_origin = ExternalRepositoryOrigin(
ManagedGrpcPythonEnvRepositoryLocationOrigin(
loadable_target_origin=loadable_target_origin
),
"doesnt_exist",
)
schedule_origin = repo_origin.get_job_origin("also_doesnt_exist")
result = sync_launch_scheduled_execution(schedule_origin)
assert isinstance(result, ScheduledExecutionFailed)
assert "doesnt_exist not found at module scope in file" in result.errors[0].to_string()
ticks = instance.get_job_ticks(schedule_origin.get_id())
assert ticks[0].status == JobTickStatus.FAILURE
assert "doesnt_exist not found at module scope in file" in ticks[0].error.message
def test_origin_ids_stable():
# This test asserts fixed schedule origin IDs to prevent any changes from
# accidentally shifting these ids that are persisted to ScheduleStorage
python_origin = ExternalJobOrigin(
ExternalRepositoryOrigin(
ManagedGrpcPythonEnvRepositoryLocationOrigin(
LoadableTargetOrigin(
executable_path="/fake/executable",
python_file="/fake/file/path",
attribute="fake_attribute",
)
),
"fake_repo",
),
"fake_schedule",
)
assert python_origin.get_id() == "eb01cc697463ba614a67567fdeaafcccc60f0fc4"
grpc_origin = ExternalJobOrigin(
ExternalRepositoryOrigin(
GrpcServerRepositoryLocationOrigin(host="fakehost", port=52618), "repo_name"
),
"fake_schedule",
)
assert grpc_origin.get_id() == "0961ecddbddfc71104adf036ebe8cd97a94dc77b"
def test_registry_multithreading():
origin = ManagedGrpcPythonEnvRepositoryLocationOrigin(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
attribute="repo",
python_file=file_relative_path(__file__,
"test_grpc_server_registry.py"),
), )
with ProcessGrpcServerRegistry(reload_interval=300,
heartbeat_ttl=600) as registry:
endpoint = registry.get_grpc_endpoint(origin)
threads = []
success_events = []
for _index in range(5):
event = threading.Event()
thread = threading.Thread(target=_registry_thread,
args=(origin, registry, endpoint, event))
threads.append(thread)
success_events.append(event)
thread.start()
for thread in threads:
thread.join()
for event in success_events:
assert event.is_set()
assert _can_connect(origin, endpoint)
registry.wait_for_processes()
assert not _can_connect(origin, endpoint)
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 sync_list_repositories_ephemeral_grpc(
executable_path,
python_file,
module_name,
working_directory,
attribute,
package_name,
):
from dagster.grpc.client import ephemeral_grpc_api_client
check.str_param(executable_path, "executable_path")
check.opt_str_param(python_file, "python_file")
check.opt_str_param(module_name, "module_name")
check.opt_str_param(working_directory, "working_directory")
check.opt_str_param(package_name, "package_name")
with ephemeral_grpc_api_client(loadable_target_origin=LoadableTargetOrigin(
executable_path=executable_path,
module_name=module_name,
python_file=python_file,
working_directory=working_directory,
attribute=attribute,
package_name=package_name,
)) as api_client:
return sync_list_repositories_grpc(api_client)
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 fake_repo_target():
return ExternalRepositoryOrigin(
ManagedGrpcPythonEnvRepositoryLocationOrigin(
LoadableTargetOrigin(executable_path=sys.executable,
module_name="fake",
attribute="fake"), ),
"fake_repo_name",
)
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 get_bar_repo_repository_location_handle():
return RepositoryLocationHandle.create_python_env_location(
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
python_file=file_relative_path(__file__, "api_tests_repo.py"),
attribute="bar_repo",
),
location_name="bar_repo_location",
)
def get_example_repository_location_handle():
loadable_target_origin = LoadableTargetOrigin(
executable_path=sys.executable, python_file=file_relative_path(__file__, "repo.py"),
)
location_name = "example_repo_location"
origin = ManagedGrpcPythonEnvRepositoryLocationOrigin(loadable_target_origin, location_name)
return RepositoryLocationHandle.create_from_repository_location_origin(origin)
def test_custom_loadable_target_origin():
# Verifies that you can swap out the LoadableTargetOrigin for the same
# repository location origin
first_loadable_target_origin = LoadableTargetOrigin(
executable_path=sys.executable,
attribute="repo",
python_file=file_relative_path(__file__,
"test_grpc_server_registry.py"),
)
second_loadable_target_origin = LoadableTargetOrigin(
executable_path=sys.executable,
attribute="other_repo",
python_file=file_relative_path(__file__,
"test_grpc_server_registry.py"),
)
origin = RegisteredRepositoryLocationOrigin("test_location")
with TestMockProcessGrpcServerRegistry() as registry:
with DynamicWorkspace(registry) as workspace:
registry.mocked_loadable_target_origin = first_loadable_target_origin
endpoint_one = registry.get_grpc_endpoint(origin)
assert registry.get_grpc_endpoint(
origin).server_id == endpoint_one.server_id
location_one = workspace.get_location(origin)
assert location_one.has_repository("repo")
# Swap in a new LoadableTargetOrigin - the same origin new returns a different
# endpoint and repository
registry.mocked_loadable_target_origin = second_loadable_target_origin
endpoint_two = registry.get_grpc_endpoint(origin)
assert endpoint_two.server_id != endpoint_one.server_id
location_two = workspace.get_location(origin)
assert location_two.has_repository("other_repo")
registry.wait_for_processes()
assert not _can_connect(origin, endpoint_one)
assert not _can_connect(origin, endpoint_two)
def test_check_run_health(kubeconfig_file):
labels = {"foo_label_key": "bar_label_value"}
# Construct a K8s run launcher in a fake k8s environment.
mock_k8s_client_batch_api = mock.Mock(spec_set=["read_namespaced_job"])
mock_k8s_client_batch_api.read_namespaced_job.side_effect = [
V1Job(status=V1JobStatus(failed=0, succeeded=0)),
V1Job(status=V1JobStatus(failed=0, succeeded=1)),
V1Job(status=V1JobStatus(failed=1, succeeded=0)),
]
k8s_run_launcher = K8sRunLauncher(
service_account_name="dagit-admin",
instance_config_map="dagster-instance",
postgres_password_secret="dagster-postgresql-secret",
dagster_home="/opt/dagster/dagster_home",
job_image="fake_job_image",
load_incluster_config=False,
kubeconfig_file=kubeconfig_file,
k8s_client_batch_api=mock_k8s_client_batch_api,
labels=labels,
)
# Create fake external pipeline.
recon_pipeline = reconstructable(fake_pipeline)
recon_repo = recon_pipeline.repository
repo_def = recon_repo.get_definition()
loadable_target_origin = LoadableTargetOrigin(python_file=__file__)
with instance_for_test() as instance:
with in_process_test_workspace(instance, loadable_target_origin) as workspace:
location = workspace.get_repository_location(workspace.repository_location_names[0])
repo_handle = RepositoryHandle(
repository_name=repo_def.name,
repository_location=location,
)
fake_external_pipeline = external_pipeline_from_recon_pipeline(
recon_pipeline,
solid_selection=None,
repository_handle=repo_handle,
)
# Launch the run in a fake Dagster instance.
pipeline_name = "demo_pipeline"
run = create_run_for_test(
instance,
pipeline_name=pipeline_name,
external_pipeline_origin=fake_external_pipeline.get_external_origin(),
pipeline_code_origin=fake_external_pipeline.get_python_origin(),
)
k8s_run_launcher.register_instance(instance)
# same order as side effects
assert k8s_run_launcher.check_run_worker_health(run).status == WorkerStatus.RUNNING
assert k8s_run_launcher.check_run_worker_health(run).status == WorkerStatus.SUCCESS
assert k8s_run_launcher.check_run_worker_health(run).status == WorkerStatus.FAILED
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_dagster_out_of_process_location():
with ManagedGrpcPythonEnvRepositoryLocationOrigin(
location_name="test_location",
loadable_target_origin=LoadableTargetOrigin(
executable_path=sys.executable,
python_file=file_relative_path(__file__, "setup.py"),
attribute="test_repo",
),
).create_test_location() as env:
assert env.get_repository("test_repo")
