def _create_tick(instance):
with ProcessGrpcServerRegistry() as grpc_server_registry:
with DynamicWorkspace(grpc_server_registry) as workspace:
list(
execute_sensor_iteration(
instance, get_default_daemon_logger("SensorDaemon"),
workspace))
def create_test_location(self):
from dagster.cli.workspace.dynamic_workspace import DynamicWorkspace
from .grpc_server_registry import ProcessGrpcServerRegistry
with ProcessGrpcServerRegistry(
reload_interval=0, heartbeat_ttl=30) as grpc_server_registry:
with DynamicWorkspace(grpc_server_registry) as workspace:
with workspace.get_location(self) as location:
yield location
def get_repository_location_from_kwargs(kwargs):
origin = get_repository_location_origin_from_kwargs(kwargs)
with ProcessGrpcServerRegistry(reload_interval=0,
heartbeat_ttl=30) as grpc_server_registry:
from dagster.cli.workspace.dynamic_workspace import DynamicWorkspace
with DynamicWorkspace(grpc_server_registry) as workspace:
with workspace.get_location(origin) as location:
yield location
def _test_launch_sensor_runs_in_subprocess(instance_ref, execution_datetime,
debug_crash_flags):
with DagsterInstance.from_ref(instance_ref) as instance:
try:
with pendulum.test(execution_datetime), ProcessGrpcServerRegistry(
) as grpc_server_registry:
with DynamicWorkspace(grpc_server_registry) as workspace:
list(
execute_sensor_iteration(
instance,
get_default_daemon_logger("SensorDaemon"),
workspace,
debug_crash_flags=debug_crash_flags,
))
finally:
cleanup_test_instance(instance)
def _test_launch_scheduled_runs_in_subprocess(instance_ref, execution_datetime, debug_crash_flags):
with DagsterInstance.from_ref(instance_ref) as instance:
try:
with ProcessGrpcServerRegistry() as grpc_server_registry:
with DynamicWorkspace(grpc_server_registry) as workspace:
with pendulum.test(execution_datetime):
list(
launch_scheduled_runs(
instance,
workspace,
logger(),
pendulum.now("UTC"),
debug_crash_flags=debug_crash_flags,
)
)
finally:
cleanup_test_instance(instance)
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_backfill_in_subprocess(instance_ref, debug_crash_flags):
execution_datetime = to_timezone(
create_pendulum_time(
year=2021,
month=2,
day=17,
),
"US/Central",
)
with DagsterInstance.from_ref(instance_ref) as instance:
try:
with pendulum.test(execution_datetime), ProcessGrpcServerRegistry(
) as grpc_server_registry:
with DynamicWorkspace(grpc_server_registry) as workspace:
list(
execute_backfill_iteration(
instance,
workspace,
get_default_daemon_logger("BackfillDaemon"),
debug_crash_flags=debug_crash_flags,
))
finally:
cleanup_test_instance(instance)
def gen_workspace(_instance):
with DynamicWorkspace(grpc_server_registry) as workspace:
yield workspace
def instance_for_context(external_repo_context, overrides=None):
with instance_for_test(overrides) as instance:
with ProcessGrpcServerRegistry() as grpc_server_registry:
with DynamicWorkspace(grpc_server_registry) as workspace:
with external_repo_context() as external_repo:
yield (instance, workspace, external_repo)
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) as registry:
with DynamicWorkspace(registry) as workspace:
endpoint_one = registry.get_grpc_endpoint(origin)
location_one = workspace.get_location(origin)
endpoint_two = registry.get_grpc_endpoint(origin)
location_two = workspace.get_location(origin)
assert endpoint_two == endpoint_one
assert location_two == location_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:
# Location manager now produces a new location as well
location_three = workspace.get_location(origin)
assert location_three != location_one
break
if time.time() - start_time > 15:
raise Exception("Server ID never changed")
time.sleep(1)
assert _can_connect(origin, endpoint_three)
# Leave workspace context, all heartbeats stop
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)