Examples with KvStateLocation org.apache.flink.runtime.query.KvStateLocation used on opensource projects

Example 1 with KvStateLocation

use of org.apache.flink.runtime.query.KvStateLocation in project flink by apache.

the class JobManagerTest method testKvStateMessages.

/**
	 * Tests that the JobManager handles {@link org.apache.flink.runtime.query.KvStateMessage}
	 * instances as expected.
	 */
@Test
public void testKvStateMessages() throws Exception {
    Deadline deadline = new FiniteDuration(100, TimeUnit.SECONDS).fromNow();
    Configuration config = new Configuration();
    config.setString(ConfigConstants.AKKA_ASK_TIMEOUT, "100ms");
    UUID leaderSessionId = null;
    ActorGateway jobManager = new AkkaActorGateway(JobManager.startJobManagerActors(config, system, TestingUtils.defaultExecutor(), TestingUtils.defaultExecutor(), TestingJobManager.class, MemoryArchivist.class)._1(), leaderSessionId);
    LeaderRetrievalService leaderRetrievalService = new StandaloneLeaderRetrievalService(AkkaUtils.getAkkaURL(system, jobManager.actor()));
    Configuration tmConfig = new Configuration();
    tmConfig.setInteger(ConfigConstants.TASK_MANAGER_MEMORY_SIZE_KEY, 4);
    tmConfig.setInteger(ConfigConstants.TASK_MANAGER_NUM_TASK_SLOTS, 8);
    ActorRef taskManager = TaskManager.startTaskManagerComponentsAndActor(tmConfig, ResourceID.generate(), system, "localhost", scala.Option.<String>empty(), scala.Option.apply(leaderRetrievalService), true, TestingTaskManager.class);
    Future<Object> registrationFuture = jobManager.ask(new NotifyWhenAtLeastNumTaskManagerAreRegistered(1), deadline.timeLeft());
    Await.ready(registrationFuture, deadline.timeLeft());
    //
    // Location lookup
    //
    LookupKvStateLocation lookupNonExistingJob = new LookupKvStateLocation(new JobID(), "any-name");
    Future<KvStateLocation> lookupFuture = jobManager.ask(lookupNonExistingJob, deadline.timeLeft()).mapTo(ClassTag$.MODULE$.<KvStateLocation>apply(KvStateLocation.class));
    try {
        Await.result(lookupFuture, deadline.timeLeft());
        fail("Did not throw expected Exception");
    } catch (IllegalStateException ignored) {
    // Expected
    }
    JobGraph jobGraph = new JobGraph("croissant");
    JobVertex jobVertex1 = new JobVertex("cappuccino");
    jobVertex1.setParallelism(4);
    jobVertex1.setMaxParallelism(16);
    jobVertex1.setInvokableClass(BlockingNoOpInvokable.class);
    JobVertex jobVertex2 = new JobVertex("americano");
    jobVertex2.setParallelism(4);
    jobVertex2.setMaxParallelism(16);
    jobVertex2.setInvokableClass(BlockingNoOpInvokable.class);
    jobGraph.addVertex(jobVertex1);
    jobGraph.addVertex(jobVertex2);
    Future<JobSubmitSuccess> submitFuture = jobManager.ask(new SubmitJob(jobGraph, ListeningBehaviour.DETACHED), deadline.timeLeft()).mapTo(ClassTag$.MODULE$.<JobSubmitSuccess>apply(JobSubmitSuccess.class));
    Await.result(submitFuture, deadline.timeLeft());
    Object lookupUnknownRegistrationName = new LookupKvStateLocation(jobGraph.getJobID(), "unknown");
    lookupFuture = jobManager.ask(lookupUnknownRegistrationName, deadline.timeLeft()).mapTo(ClassTag$.MODULE$.<KvStateLocation>apply(KvStateLocation.class));
    try {
        Await.result(lookupFuture, deadline.timeLeft());
        fail("Did not throw expected Exception");
    } catch (UnknownKvStateLocation ignored) {
    // Expected
    }
    //
    // Registration
    //
    NotifyKvStateRegistered registerNonExistingJob = new NotifyKvStateRegistered(new JobID(), new JobVertexID(), new KeyGroupRange(0, 0), "any-name", new KvStateID(), new KvStateServerAddress(InetAddress.getLocalHost(), 1233));
    jobManager.tell(registerNonExistingJob);
    LookupKvStateLocation lookupAfterRegistration = new LookupKvStateLocation(registerNonExistingJob.getJobId(), registerNonExistingJob.getRegistrationName());
    lookupFuture = jobManager.ask(lookupAfterRegistration, deadline.timeLeft()).mapTo(ClassTag$.MODULE$.<KvStateLocation>apply(KvStateLocation.class));
    try {
        Await.result(lookupFuture, deadline.timeLeft());
        fail("Did not throw expected Exception");
    } catch (IllegalStateException ignored) {
    // Expected
    }
    NotifyKvStateRegistered registerForExistingJob = new NotifyKvStateRegistered(jobGraph.getJobID(), jobVertex1.getID(), new KeyGroupRange(0, 0), "register-me", new KvStateID(), new KvStateServerAddress(InetAddress.getLocalHost(), 1293));
    jobManager.tell(registerForExistingJob);
    lookupAfterRegistration = new LookupKvStateLocation(registerForExistingJob.getJobId(), registerForExistingJob.getRegistrationName());
    lookupFuture = jobManager.ask(lookupAfterRegistration, deadline.timeLeft()).mapTo(ClassTag$.MODULE$.<KvStateLocation>apply(KvStateLocation.class));
    KvStateLocation location = Await.result(lookupFuture, deadline.timeLeft());
    assertNotNull(location);
    assertEquals(jobGraph.getJobID(), location.getJobId());
    assertEquals(jobVertex1.getID(), location.getJobVertexId());
    assertEquals(jobVertex1.getMaxParallelism(), location.getNumKeyGroups());
    assertEquals(1, location.getNumRegisteredKeyGroups());
    KeyGroupRange keyGroupRange = registerForExistingJob.getKeyGroupRange();
    assertEquals(1, keyGroupRange.getNumberOfKeyGroups());
    assertEquals(registerForExistingJob.getKvStateId(), location.getKvStateID(keyGroupRange.getStartKeyGroup()));
    assertEquals(registerForExistingJob.getKvStateServerAddress(), location.getKvStateServerAddress(keyGroupRange.getStartKeyGroup()));
    //
    // Unregistration
    //
    NotifyKvStateUnregistered unregister = new NotifyKvStateUnregistered(registerForExistingJob.getJobId(), registerForExistingJob.getJobVertexId(), registerForExistingJob.getKeyGroupRange(), registerForExistingJob.getRegistrationName());
    jobManager.tell(unregister);
    lookupFuture = jobManager.ask(lookupAfterRegistration, deadline.timeLeft()).mapTo(ClassTag$.MODULE$.<KvStateLocation>apply(KvStateLocation.class));
    try {
        Await.result(lookupFuture, deadline.timeLeft());
        fail("Did not throw expected Exception");
    } catch (UnknownKvStateLocation ignored) {
    // Expected
    }
    //
    // Duplicate registration fails task
    //
    NotifyKvStateRegistered register = new NotifyKvStateRegistered(jobGraph.getJobID(), jobVertex1.getID(), new KeyGroupRange(0, 0), "duplicate-me", new KvStateID(), new KvStateServerAddress(InetAddress.getLocalHost(), 1293));
    NotifyKvStateRegistered duplicate = new NotifyKvStateRegistered(jobGraph.getJobID(), // <--- different operator, but...
    jobVertex2.getID(), new KeyGroupRange(0, 0), // ...same name
    "duplicate-me", new KvStateID(), new KvStateServerAddress(InetAddress.getLocalHost(), 1293));
    Future<TestingJobManagerMessages.JobStatusIs> failedFuture = jobManager.ask(new NotifyWhenJobStatus(jobGraph.getJobID(), JobStatus.FAILED), deadline.timeLeft()).mapTo(ClassTag$.MODULE$.<JobStatusIs>apply(JobStatusIs.class));
    jobManager.tell(register);
    jobManager.tell(duplicate);
    // Wait for failure
    JobStatusIs jobStatus = Await.result(failedFuture, deadline.timeLeft());
    assertEquals(JobStatus.FAILED, jobStatus.state());
}

Example 2 with KvStateLocation

use of org.apache.flink.runtime.query.KvStateLocation in project flink by apache.

the class KvStateHandler method requestKvStateLocation.

public KvStateLocation requestKvStateLocation(final JobID jobId, final String registrationName) throws UnknownKvStateLocation, FlinkJobNotFoundException {
    // sanity check for the correct JobID
    if (executionGraph.getJobID().equals(jobId)) {
        if (LOG.isDebugEnabled()) {
            LOG.debug("Lookup key-value state for job {} with registration " + "name {}.", executionGraph.getJobID(), registrationName);
        }
        final KvStateLocationRegistry registry = executionGraph.getKvStateLocationRegistry();
        final KvStateLocation location = registry.getKvStateLocation(registrationName);
        if (location != null) {
            return location;
        } else {
            throw new UnknownKvStateLocation(registrationName);
        }
    } else {
        if (LOG.isDebugEnabled()) {
            LOG.debug("Request of key-value state location for unknown job {} received.", jobId);
        }
        throw new FlinkJobNotFoundException(jobId);
    }
}

Example 3 with KvStateLocation

use of org.apache.flink.runtime.query.KvStateLocation in project flink by apache.

the class KvStateClientProxyHandler method getKvStateLookupInfo.

/**
 * Lookup the {@link KvStateLocation} for the given job and queryable state name.
 *
 * <p>The job manager will be queried for the location only if forced or no cached location can
 * be found. There are no guarantees about
 *
 * @param jobId JobID the state instance belongs to.
 * @param queryableStateName Name under which the state instance has been published.
 * @param forceUpdate Flag to indicate whether to force a update via the lookup service.
 * @return Future holding the KvStateLocation
 */
private CompletableFuture<KvStateLocation> getKvStateLookupInfo(final JobID jobId, final String queryableStateName, final boolean forceUpdate) {
    final Tuple2<JobID, String> cacheKey = new Tuple2<>(jobId, queryableStateName);
    final CompletableFuture<KvStateLocation> cachedFuture = lookupCache.get(cacheKey);
    if (!forceUpdate && cachedFuture != null && !cachedFuture.isCompletedExceptionally()) {
        LOG.debug("Retrieving location for state={} of job={} from the cache.", queryableStateName, jobId);
        return cachedFuture;
    }
    final KvStateLocationOracle kvStateLocationOracle = proxy.getKvStateLocationOracle(jobId);
    if (kvStateLocationOracle != null) {
        LOG.debug("Retrieving location for state={} of job={} from the key-value state location oracle.", queryableStateName, jobId);
        final CompletableFuture<KvStateLocation> location = new CompletableFuture<>();
        lookupCache.put(cacheKey, location);
        kvStateLocationOracle.requestKvStateLocation(jobId, queryableStateName).whenComplete((KvStateLocation kvStateLocation, Throwable throwable) -> {
            if (throwable != null) {
                if (ExceptionUtils.stripCompletionException(throwable) instanceof FlinkJobNotFoundException) {
                    // if the jobId was wrong, remove the entry from the cache.
                    lookupCache.remove(cacheKey);
                }
                location.completeExceptionally(throwable);
            } else {
                location.complete(kvStateLocation);
            }
        });
        return location;
    } else {
        return FutureUtils.completedExceptionally(new UnknownLocationException("Could not retrieve location of state=" + queryableStateName + " of job=" + jobId + ". Potential reasons are: i) the state is not ready, or ii) the job does not exist."));
    }
}

Example 4 with KvStateLocation

use of org.apache.flink.runtime.query.KvStateLocation in project flink by apache.

the class JobMaster method lookupKvStateLocation.

@RpcMethod
public KvStateLocation lookupKvStateLocation(final String registrationName) throws Exception {
    if (log.isDebugEnabled()) {
        log.debug("Lookup key-value state for job {} with registration " + "name {}.", jobGraph.getJobID(), registrationName);
    }
    final KvStateLocationRegistry registry = executionGraph.getKvStateLocationRegistry();
    final KvStateLocation location = registry.getKvStateLocation(registrationName);
    if (location != null) {
        return location;
    } else {
        throw new UnknownKvStateLocation(registrationName);
    }
}

Example 5 with KvStateLocation

use of org.apache.flink.runtime.query.KvStateLocation in project flink by apache.

the class JobMasterQueryableStateTest method testRegisterKvState.

@Test
public void testRegisterKvState() throws Exception {
    final JobMaster jobMaster = new JobMasterBuilder(JOB_GRAPH, rpcService).createJobMaster();
    jobMaster.start();
    final JobMasterGateway jobMasterGateway = jobMaster.getSelfGateway(JobMasterGateway.class);
    registerSlotsRequiredForJobExecution(jobMasterGateway, JOB_GRAPH.getJobID());
    try {
        final String registrationName = "register-me";
        final KvStateID kvStateID = new KvStateID();
        final KeyGroupRange keyGroupRange = new KeyGroupRange(0, 0);
        final InetSocketAddress address = new InetSocketAddress(InetAddress.getLocalHost(), 1029);
        jobMasterGateway.notifyKvStateRegistered(JOB_GRAPH.getJobID(), JOB_VERTEX_1.getID(), keyGroupRange, registrationName, kvStateID, address).get();
        final KvStateLocation location = jobMasterGateway.requestKvStateLocation(JOB_GRAPH.getJobID(), registrationName).get();
        assertEquals(JOB_GRAPH.getJobID(), location.getJobId());
        assertEquals(JOB_VERTEX_1.getID(), location.getJobVertexId());
        assertEquals(JOB_VERTEX_1.getMaxParallelism(), location.getNumKeyGroups());
        assertEquals(1, location.getNumRegisteredKeyGroups());
        assertEquals(1, keyGroupRange.getNumberOfKeyGroups());
        assertEquals(kvStateID, location.getKvStateID(keyGroupRange.getStartKeyGroup()));
        assertEquals(address, location.getKvStateServerAddress(keyGroupRange.getStartKeyGroup()));
    } finally {
        RpcUtils.terminateRpcEndpoint(jobMaster, testingTimeout);
    }
}