def test_user_defined_k8s_config_in_run_tags(kubeconfig_file): # Construct a K8s run launcher in a fake k8s environment. mock_k8s_client_batch_api = mock.MagicMock() celery_k8s_run_launcher = CeleryK8sRunLauncher( instance_config_map="dagster-instance", postgres_password_secret="dagster-postgresql-secret", dagster_home="/opt/dagster/dagster_home", load_incluster_config=False, kubeconfig_file=kubeconfig_file, k8s_client_batch_api=mock_k8s_client_batch_api, ) # Construct Dagster run tags with user defined k8s config. expected_resources = { "requests": {"cpu": "250m", "memory": "64Mi"}, "limits": {"cpu": "500m", "memory": "2560Mi"}, } user_defined_k8s_config = UserDefinedDagsterK8sConfig( container_config={"resources": expected_resources}, ) user_defined_k8s_config_json = json.dumps(user_defined_k8s_config.to_dict()) tags = {"dagster-k8s/config": user_defined_k8s_config_json} # Create fake external pipeline. recon_pipeline = reconstructable(fake_pipeline) recon_repo = recon_pipeline.repository location_origin = InProcessRepositoryLocationOrigin(recon_repo) location_handle = location_origin.create_handle() repo_def = recon_repo.get_definition() repo_handle = RepositoryHandle( repository_name=repo_def.name, repository_location_handle=location_handle, ) 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. with instance_for_test() as instance: celery_k8s_run_launcher.initialize(instance) pipeline_name = "demo_pipeline" run_config = {"execution": {"celery-k8s": {"config": {"job_image": "fake-image-name"}}}} run = create_run_for_test( instance, pipeline_name=pipeline_name, run_config=run_config, tags=tags, ) celery_k8s_run_launcher.launch_run(instance, run, fake_external_pipeline) # Check that user defined k8s config was passed down to the k8s job. 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" job_resources = kwargs["body"].spec.template.spec.containers[0].resources assert job_resources == expected_resources
def test_no_postgres(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", 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, ) # Create fake external pipeline. recon_pipeline = reconstructable(fake_pipeline) recon_repo = recon_pipeline.repository repo_def = recon_repo.get_definition() location_origin = InProcessRepositoryLocationOrigin(recon_repo) with location_origin.create_location() as location: 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. with instance_for_test() as instance: pipeline_name = "demo_pipeline" run = create_run_for_test(instance, pipeline_name=pipeline_name) k8s_run_launcher.register_instance(instance) run = k8s_run_launcher.launch_run(run, fake_external_pipeline) updated_run = instance.get_run_by_id(run.run_id) assert updated_run.tags[DOCKER_IMAGE_TAG] == "fake_job_image" # Check that user defined k8s config was passed down to the k8s job. 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" assert DAGSTER_PG_PASSWORD_ENV_VAR not in [ env.name for env in kwargs["body"].spec.template.spec.containers[0].env ]
def _get_unloadable_sensor_origin(): working_directory = os.path.dirname(__file__) recon_repo = ReconstructableRepository.for_file(__file__, "doesnt_exist", working_directory) return ExternalRepositoryOrigin( InProcessRepositoryLocationOrigin(recon_repo), "fake_repository").get_job_origin("doesnt_exist")
def get_external_origin(self): """ Hack! Inject origin that the k8s images will use. The BK image uses a different directory structure (/workdir/python_modules/dagster-test/dagster_test/test_project) than the images inside the kind cluster (/dagster_test/test_project). As a result the normal origin won't work, we need to inject this one. """ return ExternalJobOrigin( external_repository_origin=ExternalRepositoryOrigin( repository_location_origin=InProcessRepositoryLocationOrigin( recon_repo=ReconstructableRepository( pointer=FileCodePointer( python_file= "/dagster_test/test_project/test_pipelines/repo.py", fn_name="define_demo_execution_repo", ), container_image=self._container_image, executable_path="python", entry_point=DEFAULT_DAGSTER_ENTRY_POINT, )), repository_name="demo_execution_repo", ), job_name=self.name, )
def get_test_project_external_pipeline(pipeline_name): return (InProcessRepositoryLocationOrigin( ReconstructableRepository.for_file( file_relative_path(__file__, "test_pipelines/repo.py"), "define_demo_execution_repo", )).create_handle().create_location().get_repository( "demo_execution_repo").get_full_external_pipeline(pipeline_name))
def _repo_location_origins_from_config(workspace_config, yaml_path): ensure_workspace_config(workspace_config, yaml_path) if "repository" in workspace_config: warnings.warn( # link to docs once they exist "You are using the legacy repository yaml format. Please update your file " "to abide by the new workspace file format." ) origin = InProcessRepositoryLocationOrigin( ReconstructableRepository.from_legacy_repository_yaml(yaml_path) ) return {origin.location_name: origin} location_origins = OrderedDict() for location_config in workspace_config["load_from"]: origin = _location_origin_from_location_config(location_config, yaml_path) check.invariant( location_origins.get(origin.location_name) is None, 'Cannot have multiple locations with the same name, got multiple "{name}"'.format( name=origin.location_name, ), ) location_origins[origin.location_name] = origin return location_origins
def legacy_get_bar_repo_handle(): recon_repo = ReconstructableRepository.from_legacy_repository_yaml( file_relative_path(__file__, "legacy_repository_file.yaml")) return (RepositoryLocation.from_handle( RepositoryLocationHandle.create_from_repository_location_origin( InProcessRepositoryLocationOrigin(recon_repo))).get_repository( "bar_repo").handle)
def get_test_project_external_repo(container_image=None): with InProcessRepositoryLocationOrigin( ReconstructableRepository.for_file( file_relative_path(__file__, "test_pipelines/repo.py"), "define_demo_execution_repo", container_image=container_image, )).create_location() as location: yield location.get_repository("demo_execution_repo")
def get_test_project_external_repo(container_image=None): return RepositoryLocation.from_handle( RepositoryLocationHandle.create_from_repository_location_origin( InProcessRepositoryLocationOrigin( ReconstructableRepository.for_file( file_relative_path(__file__, "test_pipelines/repo.py"), "define_demo_execution_repo", container_image=container_image, )))).get_repository("demo_execution_repo")
def define_in_process_context(python_file, fn_name, instance): check.inst_param(instance, "instance", DagsterInstance) return WorkspaceProcessContext( workspace=Workspace([ InProcessRepositoryLocationOrigin( ReconstructableRepository.for_file(python_file, fn_name)) ]), instance=instance, ).create_request_context()
def get_test_project_external_pipeline(pipeline_name): with InProcessRepositoryLocationOrigin( LoadableTargetOrigin( executable_path=sys.executable, python_file=file_relative_path(__file__, "test_pipelines/repo.py"), attribute="define_demo_execution_repo", )).create_location() as location: yield location.get_repository( "demo_execution_repo").get_full_external_pipeline(pipeline_name)
def _get_unloadable_sensor_origin(name): working_directory = os.path.dirname(__file__) loadable_target_origin = LoadableTargetOrigin( executable_path=sys.executable, python_file=__file__, working_directory=working_directory, ) return ExternalRepositoryOrigin( InProcessRepositoryLocationOrigin(loadable_target_origin), "fake_repository").get_instigator_origin(name)
def _unloadable_partition_set_origin(): working_directory = os.path.dirname(__file__) return ExternalRepositoryOrigin( InProcessRepositoryLocationOrigin( LoadableTargetOrigin( executable_path=sys.executable, python_file=__file__, working_directory=working_directory, )), "fake_repository", ).get_partition_set_origin("doesnt_exist")
def _mgr_fn(recon_repo): check.inst_param(recon_repo, "recon_repo", ReconstructableRepository) with Workspace([InProcessRepositoryLocationOrigin(recon_repo)]) as workspace: yield workspace
def test_user_defined_k8s_config_in_run_tags(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, ) # Construct Dagster run tags with user defined k8s config. expected_resources = { "requests": { "cpu": "250m", "memory": "64Mi" }, "limits": { "cpu": "500m", "memory": "2560Mi" }, } user_defined_k8s_config = UserDefinedDagsterK8sConfig( container_config={"resources": expected_resources}, ) user_defined_k8s_config_json = json.dumps( user_defined_k8s_config.to_dict()) tags = {"dagster-k8s/config": user_defined_k8s_config_json} # Create fake external pipeline. recon_pipeline = reconstructable(fake_pipeline) recon_repo = recon_pipeline.repository repo_def = recon_repo.get_definition() location_origin = InProcessRepositoryLocationOrigin(recon_repo) with location_origin.create_location() as location: 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. with instance_for_test() as instance: pipeline_name = "demo_pipeline" run = create_run_for_test(instance, pipeline_name=pipeline_name, tags=tags) k8s_run_launcher.register_instance(instance) run = k8s_run_launcher.launch_run(run, fake_external_pipeline) updated_run = instance.get_run_by_id(run.run_id) assert updated_run.tags[DOCKER_IMAGE_TAG] == "fake_job_image" # Check that user defined k8s config was passed down to the k8s job. 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" job_resources = kwargs["body"].spec.template.spec.containers[ 0].resources assert job_resources == expected_resources assert DAGSTER_PG_PASSWORD_ENV_VAR in [ env.name for env in kwargs["body"].spec.template.spec.containers[0].env ]