FlyteAdminClientRunner implements FlyteRunner { @Override public String createExecution(final RunState runState, final String execName, final FlyteExecConf flyteExecConf, final Map<String, String> annotations) throws CreateExecutionException { requireNonNull(runState, "runState"); requireNonNull(execName, "name"); requireNonNull(flyteExecConf, "flyteExecConf"); requireNonNull(annotations, "annotations"); final var launchPlanIdentifier = flyteExecConf.referenceId(); final var execMode = runState.data().trigger() .map(this::toFlyteExecutionMode) .filter(mode -> mode != ExecutionMode.UNRECOGNIZED) .orElseThrow(() -> new CreateExecutionException("Missing trigger or unknown in StateData: " + runState.data())); final var parameter = runState.workflowInstance().parameter(); try { flyteAdminClient.createExecution( launchPlanIdentifier.project(), launchPlanIdentifier.domain(), execName, IdentifierOuterClass.Identifier.newBuilder() .setName(launchPlanIdentifier.name()) .setProject(launchPlanIdentifier.project()) .setDomain(launchPlanIdentifier.domain()) .setResourceType(IdentifierOuterClass.ResourceType.LAUNCH_PLAN) .setVersion(launchPlanIdentifier.version()) .build(), execMode, annotations, parameter); return runnerId; } catch (StatusRuntimeException e) { switch (e.getStatus().getCode()) { case ALREADY_EXISTS: return runnerId; case NOT_FOUND: throw new LaunchPlanNotFound(flyteExecConf, e); } throw new CreateExecutionException(flyteExecConf, e); } catch (Exception e) { throw new CreateExecutionException(flyteExecConf, e); } } @VisibleForTesting  FlyteAdminClientRunner(final String runnerId,
                         final FlyteAdminClient flyteAdminClient,
                         final StateManager stateManager,
                         final Duration terminateDanglingFlyteExecInterval,
                         final ScheduledExecutorService scheduledExecutor,
                         final Time time);   FlyteAdminClientRunner(final String runnerId,
                         final FlyteAdminClient flyteAdminClient,
                         final StateManager stateManager); @Override String createExecution(final RunState runState, final String execName, final FlyteExecConf flyteExecConf,
                                final Map<String, String> annotations); @Override void terminateExecution(RunState runState, FlyteExecutionId flyteExecutionId); @Override void poll(FlyteExecutionId flyteExecutionId, RunState runState); @Override void close();  }
@Test public void testCreateExecution() throws FlyteRunner.CreateExecutionException { final var execName = "test-create-execution"; when(flyteAdminClient.createExecution(any(), any(), any(), any(), any(), any(), any())).thenReturn( ExecutionOuterClass.ExecutionCreateResponse .newBuilder() .setId(IdentifierOuterClass.WorkflowExecutionIdentifier .newBuilder() .setProject(LAUNCH_PLAN_IDENTIFIER.project()) .setDomain(LAUNCH_PLAN_IDENTIFIER.domain()) .setName(execName) .build()) .build()); var runnerId = flyteRunner.createExecution(RUN_STATE, execName, FLYTE_EXEC_CONF, ANNOTATIONS); assertThat(runnerId, is(RUNNER_ID)); verify(flyteAdminClient).createExecution( LAUNCH_PLAN_IDENTIFIER.project(), LAUNCH_PLAN_IDENTIFIER.domain(), execName, toProto(LAUNCH_PLAN_IDENTIFIER), SCHEDULED, ANNOTATIONS, RUN_STATE.workflowInstance().parameter()); }
@Test @Parameters(method = "styxTriggerToFlyteExecMode") public void testStyxTriggerTranslateToFlyteExecutionMode(Trigger styxTrigger, ExecutionMode flyteExecMode) throws FlyteRunner.CreateExecutionException { final RunState runState = runState(styxTrigger); flyteRunner.createExecution(runState, "test-create-execution", FLYTE_EXEC_CONF, ANNOTATIONS); verify(flyteAdminClient).createExecution(any(), any(), any(), any(), Mockito.eq(flyteExecMode), any(), any()); }
@Test @Parameters(method = "invalidsStyxTriggerToFlyteExecMode") public void testCreateExecutionThrowsExceptionForInvalidExecutionMode(Trigger styxTrigger) { var exception = assertThrows( FlyteRunner.CreateExecutionException.class, () -> flyteRunner.createExecution(runState(styxTrigger), "test-create-execution", FLYTE_EXEC_CONF, ANNOTATIONS) ); assertThat(exception.getMessage(), containsString("Missing trigger or unknown in StateData")); }
@Test public void testThrowsFlyteLaunchPlanNotFound() { doThrow(new StatusRuntimeException(Status.NOT_FOUND)) .when(flyteAdminClient).createExecution(any(), any(), any(), any(), any(), any(), any()); assertThrows( FlyteRunner.LaunchPlanNotFound.class, () -> flyteRunner.createExecution(RUN_STATE, "exec", FLYTE_EXEC_CONF, ANNOTATIONS)); }
@Test public void testCreateExecutionForAlreadyExistsException() throws FlyteRunner.CreateExecutionException { doThrow(new StatusRuntimeException(Status.ALREADY_EXISTS)) .when(flyteAdminClient).createExecution(any(), any(), any(), any(), any(), any(), any()); var runnerId = flyteRunner.createExecution(RUN_STATE, "exec", FLYTE_EXEC_CONF, ANNOTATIONS); assertThat(runnerId, is(RUNNER_ID)); verify(flyteAdminClient).createExecution( LAUNCH_PLAN_IDENTIFIER.project(), LAUNCH_PLAN_IDENTIFIER.domain(), "exec", toProto(LAUNCH_PLAN_IDENTIFIER), SCHEDULED, ANNOTATIONS, RUN_STATE.workflowInstance().parameter()); }
@Test public void testThrowsCreateExecutionExceptionForOtherCode() { doThrow(new StatusRuntimeException(Status.INTERNAL)) .when(flyteAdminClient).createExecution(any(), any(), any(), any(), any(), any(), any()); assertThrows( FlyteRunner.CreateExecutionException.class, () -> flyteRunner.createExecution(RUN_STATE, "exec", FLYTE_EXEC_CONF, ANNOTATIONS)); }
@Test public void testThrowsCreateExecutionExceptionForUnknownException() { doThrow(new IllegalStateException("test")) .when(flyteAdminClient).createExecution(any(), any(), any(), any(), any(), any(), any()); assertThrows( FlyteRunner.CreateExecutionException.class, () -> flyteRunner.createExecution(RUN_STATE, "exec", FLYTE_EXEC_CONF, ANNOTATIONS)); }