Skip to content

[CELEBORN-2394] Add REST v1 API to unregister shuffles #3770

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Kalvin2077 wants to merge 4 commits into from
Closed

[CELEBORN-2394] Add REST v1 API to unregister shuffles #3770

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b0c27cf
feat: Add REST v1 API to unregister a shuffle
Kalvin2077 Jul 28, 2026
322b2f3
fix: legacy test use incorrect type.
Kalvin2077 Jul 28, 2026
643a2ba
feat: Support unregistering multiple shuffles
Kalvin2077 Jul 29, 2026
65f6d2f
refactor: Reuse shuffle CLI options
Kalvin2077 Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions cli/src/main/scala/org/apache/celeborn/cli/master/MasterOptions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,11 @@ final class MasterOptions {
@Option(names = Array("--show-cluster-shuffles"), description = Array("Show cluster shuffles"))
private[master] var showClusterShuffles: Boolean = _

@Option(
Comment thread
SteNicholas marked this conversation as resolved.
names = Array("--unregister-shuffles"),
description = Array("Unregister shuffles from the service"))
private[master] var unregisterShuffles: Boolean = _

@Option(names = Array("--exclude-worker"), description = Array("Exclude workers by ID"))
private[master] var excludeWorkers: Boolean = _

Expand Down
4 changes: 3 additions & 1 deletion cli/src/main/scala/org/apache/celeborn/cli/master/MasterSubcommand.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ trait MasterSubcommand extends BaseCommand {
private[master] var masterOptions: MasterOptions = _

@ArgGroup(exclusive = false)
private[master] var reviseLostShuffleOptions: ReviseLostShuffleOptions = _
private[master] var shuffleOptions: ShuffleOptions = _

@Mixin
private[master] var commonOptions: CommonOptions = _
Expand Down Expand Up @@ -77,6 +77,8 @@ trait MasterSubcommand extends BaseCommand {

private[master] def runShowClusterShuffles: ShufflesResponse

private[master] def runUnregisterShuffles: HandleResponse

private[master] def runExcludeWorkers: HandleResponse

private[master] def runRemoveExcludedWorkers: HandleResponse
Expand Down
35 changes: 26 additions & 9 deletions cli/src/main/scala/org/apache/celeborn/cli/master/MasterSubcommandImpl.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ class MasterSubcommandImpl extends MasterSubcommand {
if (masterOptions.showClusterApps) log(runShowClusterApps)
if (masterOptions.showClusterAppsInfo) log(runShowClusterAppsInfo)
if (masterOptions.showClusterShuffles) log(runShowClusterShuffles)
if (masterOptions.unregisterShuffles) log(runUnregisterShuffles)
if (masterOptions.excludeWorkers) log(runExcludeWorkers)
if (masterOptions.removeExcludedWorkers) log(runRemoveExcludedWorkers)
if (masterOptions.removeWorkersUnavailableInfo) log(runRemoveWorkersUnavailableInfo)
Expand Down Expand Up @@ -77,6 +78,20 @@ class MasterSubcommandImpl extends MasterSubcommand {
private[master] def runShowClusterShuffles: ShufflesResponse =
shuffleApi.getShuffles(commonOptions.getAuthHeader)

private[master] def runUnregisterShuffles: HandleResponse = {
val (appId, shuffleIds) = getSingleAppShuffleIds
if (shuffleIds.asScala.exists(_ < 0)) {
throw new ParameterException(
spec.commandLine(),
"Shuffle ids must be nonnegative.")
}

val request = new UnregisterShufflesRequest()
.appId(appId)
.shuffleIds(shuffleIds)
shuffleApi.unregisterShuffles(request, commonOptions.getAuthHeader)
}

private[master] def runExcludeWorkers: HandleResponse = {
val workerIds = getWorkerIds
val excludeWorkerRequest = new ExcludeWorkerRequest().add(workerIds)
Expand Down Expand Up @@ -251,24 +266,26 @@ class MasterSubcommandImpl extends MasterSubcommand {
private[master] def runShowContainerInfo: ContainerInfo =
defaultApi.getContainerInfo(commonOptions.getAuthHeader)

override private[master] def reviseLostShuffles: HandleResponse = {
if (StringUtils.isAnyBlank(commonOptions.apps, reviseLostShuffleOptions.shuffleIds)) {
private def getSingleAppShuffleIds: (String, util.List[Integer]) = {
val appId = commonOptions.apps
val shuffleIds = Option(shuffleOptions).map(_.shuffleIds).orNull
if (StringUtils.isBlank(appId) || shuffleIds == null || shuffleIds.isEmpty) {
throw new ParameterException(
spec.commandLine(),
"Application id and Shuffle ids must be provided for this command.")
"Application id and shuffle ids must be provided for this command.")
}

val app = commonOptions.apps
if (app.contains(",")) {
if (appId.contains(",")) {
throw new ParameterException(
spec.commandLine(),
"Only one application id can be provided for this command.")
}
(appId, shuffleIds)
}

val shuffleIds = util.Arrays.asList[Integer](
reviseLostShuffleOptions.shuffleIds.split(",").map(Integer.valueOf): _*)
override private[master] def reviseLostShuffles: HandleResponse = {
val (appId, shuffleIds) = getSingleAppShuffleIds
val request =
new ReviseLostShufflesRequest().appId(app).shuffleIds(shuffleIds)
new ReviseLostShufflesRequest().appId(appId).shuffleIds(shuffleIds)
applicationApi.reviseLostShuffles(request, commonOptions.getAuthHeader)
}

Expand Down
8 changes: 6 additions & 2 deletions ...cli/master/ReviseLostShuffleOptions.scala → .../celeborn/cli/master/ShuffleOptions.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,13 +17,17 @@

package org.apache.celeborn.cli.master

import java.util

import picocli.CommandLine.Option

final class ReviseLostShuffleOptions {
final class ShuffleOptions {

@Option(
names = Array("--shuffleIds"),
paramLabel = "shuffleId",
split = ",",
description = Array("The shuffle ids to manipulate."))
private[master] var shuffleIds: String = _
private[master] var shuffleIds: util.List[Integer] = _

}
26 changes: 26 additions & 0 deletions cli/src/test/scala/org/apache/celeborn/cli/TestCelebornCliCommands.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,32 @@ class TestCelebornCliCommands extends CelebornFunSuite with MiniClusterFeature {
captureOutputAndValidateResponse(args, "ShufflesResponse")
}

test("master --unregister-shuffles") {
val args = prepareMasterArgs() ++ Array(
"--unregister-shuffles",
"--apps",
"app1",
"--shuffleIds",
"1,2")
captureOutputAndValidateResponse(args, "Unregistered shuffles app1-1, app1-2.")
}

test("master --unregister-shuffles validates inputs") {
Seq(
Array("--unregister-shuffles", "--shuffleIds", "1,2") ->
"Application id and shuffle ids must be provided",
Array("--unregister-shuffles", "--apps", "app1") ->
"Application id and shuffle ids must be provided",
Array("--unregister-shuffles", "--apps", "app1,app2", "--shuffleIds", "1,2") ->
"Only one application id can be provided",
Array("--unregister-shuffles", "--apps", "app1", "--shuffleIds", "1,invalid") ->
"Invalid value for option '--shuffleIds'",
Array("--unregister-shuffles", "--apps", "app1", "--shuffleIds", "1,-1") ->
"Shuffle ids must be nonnegative").foreach { case (command, expectedError) =>
captureErrorAndValidateResponse(prepareMasterArgs() ++ command, expectedError)
}
}

test("master --show-worker-event-info") {
val args = prepareMasterArgs() :+ "--show-worker-event-info"
captureOutputAndValidateResponse(args, "WorkerEventsResponse")
Expand Down
12 changes: 7 additions & 5 deletions docs/celeborn_cli.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,9 +86,10 @@ Usage: celeborn-cli master [-hV] [--apps=appId] [--auth-header=authHeader]
v1,k2:v2,k3:v3...] [--worker-ids=w1,w2,w3...]
(--show-masters-info | --show-cluster-apps |
--show-cluster-apps-info | --show-cluster-shuffles |
--exclude-worker | --remove-excluded-worker |
--send-worker-event=IMMEDIATELY | DECOMMISSION |
DECOMMISSION_THEN_IDLE | GRACEFUL | RECOMMISSION |
--unregister-shuffles | --exclude-worker |
--remove-excluded-worker |
--send-worker-event=IMMEDIATELY | DECOMMISSION |
DECOMMISSION_THEN_IDLE | GRACEFUL | RECOMMISSION |
NONE | --show-worker-event-info |
--show-lost-workers | --show-excluded-workers |
--show-manual-excluded-workers |
Expand All @@ -104,7 +105,7 @@ Usage: celeborn-cli master [-hV] [--apps=appId] [--auth-header=authHeader]
--revise-lost-shuffles | --delete-apps |
--update-interruption-notices=workerId1=timestamp,
workerId2=timestamp,workerId3=timestamp)
[[--shuffleIds=<shuffleIds>]]
[[--shuffleIds=shuffleId[,shuffleId...]]...]
--add-cluster-alias=alias
Add alias to use in the cli for the given set of
masters
Expand Down Expand Up @@ -168,8 +169,9 @@ Usage: celeborn-cli master [-hV] [--apps=appId] [--auth-header=authHeader]
--show-workers Show registered workers
--show-workers-topology
Show registered workers topology
--shuffleIds=<shuffleIds>
--shuffleIds=shuffleId[,shuffleId...]
The shuffle ids to manipulate.
--unregister-shuffles Unregister shuffles from the service
--update-interruption-notices=workerId1=timestamp,workerId2=timestamp,
workerId3=timestamp
Update interruption notices of workers.
Expand Down
16 changes: 16 additions & 0 deletions master/src/main/scala/org/apache/celeborn/service/deploy/master/Master.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1461,6 +1461,22 @@ private[celeborn] class Master(
sb.toString()
}

def unregisterShuffles(
applicationId: String,
shuffleIds: util.List[Integer]): HandleResponse = {
val shuffleKeys =
shuffleIds.asScala.map(Utils.makeShuffleKey(applicationId, _)).mkString(", ")
val response = self.askSync[PbBatchUnregisterShuffleResponse](
BatchUnregisterShuffles(applicationId, shuffleIds, MasterClient.genRequestId()))
val status = StatusCode.fromValue(response.getStatus)
val success = status == StatusCode.SUCCESS
if (success) {
(success, s"Unregistered shuffles $shuffleKeys.")
} else {
(success, s"Failed to unregister shuffles $shuffleKeys: $status.")
}
}

override def exclude(
addWorkers: Seq[WorkerInfo],
removeWorkers: Seq[WorkerInfo]): HandleResponse = {
Expand Down
34 changes: 31 additions & 3 deletions ...rc/main/scala/org/apache/celeborn/service/deploy/master/http/api/v1/ShuffleResource.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
package org.apache.celeborn.service.deploy.master.http.api.v1

import java.util
import javax.ws.rs.{Consumes, GET, Produces}
import javax.ws.rs.{BadRequestException, Consumes, GET, Path, POST, Produces}
import javax.ws.rs.core.MediaType

import scala.collection.JavaConverters._
Expand All @@ -28,15 +28,17 @@ import io.swagger.v3.oas.annotations.media.{Content, Schema}
import io.swagger.v3.oas.annotations.responses.ApiResponse
import io.swagger.v3.oas.annotations.tags.Tag

import org.apache.celeborn.rest.v1.model.ShufflesResponse
import org.apache.celeborn.rest.v1.model.{HandleResponse, ShufflesResponse, UnregisterShufflesRequest}
import org.apache.celeborn.server.common.http.api.ApiRequestContext
import org.apache.celeborn.service.deploy.master.Master
import org.apache.celeborn.service.deploy.master.http.api.MasterHttpResourceUtils.ensureMasterIsLeader

@Tag(name = "Shuffle")
@Produces(Array(MediaType.APPLICATION_JSON))
@Consumes(Array(MediaType.APPLICATION_JSON))
class ShuffleResource extends ApiRequestContext {
private def statusSystem = httpService.asInstanceOf[Master].statusSystem
private def master = httpService.asInstanceOf[Master]
private def statusSystem = master.statusSystem

@Operation(description =
"List all running shuffle keys of the service. It will return all running shuffle's key of the cluster.")
Expand All @@ -56,4 +58,30 @@ class ShuffleResource extends ApiRequestContext {
}
new ShufflesResponse().shuffleIds(shuffles)
}

@Operation(description = "Unregister shuffles from the service.")
@ApiResponse(
responseCode = "200",
content = Array(new Content(
mediaType = MediaType.APPLICATION_JSON,
schema = new Schema(implementation = classOf[HandleResponse]))))
@POST
@Path("/unregister")
def unregisterShuffles(request: UnregisterShufflesRequest): HandleResponse =
ensureMasterIsLeader(master) {
if (request == null) {
throw new BadRequestException("The unregister shuffles request is required.")
}
val appId = normalizeParam(request.getAppId)
val shuffleIds = request.getShuffleIds
if (appId.isEmpty ||
shuffleIds == null ||
shuffleIds.isEmpty ||
shuffleIds.asScala.exists(shuffleId => shuffleId == null || shuffleId < 0)) {
throw new BadRequestException(
s"appId(${request.getAppId}) is required and shuffleIds($shuffleIds) must be a nonempty list of nonnegative ids.")
}
val (success, message) = master.unregisterShuffles(appId, shuffleIds)
new HandleResponse().success(success).message(message)
}
}
91 changes: 88 additions & 3 deletions ...cala/org/apache/celeborn/service/deploy/master/http/api/v1/ApiV1MasterResourceSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,13 @@

package org.apache.celeborn.service.deploy.master.http.api.v1

import java.util.Collections
import java.util.{Arrays, Collections}
import javax.servlet.http.HttpServletResponse
import javax.ws.rs.client.Entity
import javax.ws.rs.core.MediaType

import org.apache.celeborn.rest.v1.model.{ApplicationsResponse, ExcludeWorkerRequest, HandleResponse, HostnamesResponse, RemoveWorkersUnavailableInfoRequest, SendWorkerEventRequest, ShufflesResponse, TopologyResponse, WorkerEventsResponse, WorkerId, WorkersResponse}
import org.apache.celeborn.common.util.Utils
import org.apache.celeborn.rest.v1.model.{ApplicationsHeartbeatResponse, ExcludeWorkerRequest, HandleResponse, HostnamesResponse, RemoveWorkersUnavailableInfoRequest, SendWorkerEventRequest, ShufflesResponse, TopologyResponse, UnregisterShufflesRequest, WorkerEventsResponse, WorkerId, WorkersResponse}
import org.apache.celeborn.server.common.HttpService
import org.apache.celeborn.server.common.http.api.v1.ApiV1BaseResourceSuite
import org.apache.celeborn.service.deploy.master.{Master, MasterClusterFeature}
Expand All @@ -48,10 +49,94 @@ class ApiV1MasterResourceSuite extends ApiV1BaseResourceSuite with MasterCluster
assert(response.readEntity(classOf[ShufflesResponse]).getShuffleIds.isEmpty)
}

test("unregister shuffles preserves other shuffles and is idempotent") {
val appId = "unregister-shuffles-app"
val shuffleIds = Arrays.asList[Integer](0, 1)
val remainingShuffleId = 2
shuffleIds.forEach { shuffleId =>
master.statusSystem.updateRequestSlotsMeta(
Utils.makeShuffleKey(appId, shuffleId),
null,
Collections.emptyMap[String, java.util.Map[String, Integer]]())
}
master.statusSystem.updateRequestSlotsMeta(
Utils.makeShuffleKey(appId, remainingShuffleId),
null,
Collections.emptyMap[String, java.util.Map[String, Integer]]())
try {
shuffleIds.forEach { shuffleId =>
assert(master.statusSystem.registeredAppAndShuffles.get(appId).contains(shuffleId))
}
assert(master.statusSystem.registeredAppAndShuffles.get(appId).contains(remainingShuffleId))

val request = new UnregisterShufflesRequest().appId(appId).shuffleIds(shuffleIds)
var response =
webTarget.path("shuffles/unregister").request(MediaType.APPLICATION_JSON).post(
Entity.entity(request, MediaType.APPLICATION_JSON))
assert(HttpServletResponse.SC_OK == response.getStatus)
val handleResponse = response.readEntity(classOf[HandleResponse])
assert(handleResponse.getSuccess)
shuffleIds.forEach { shuffleId =>
assert(handleResponse.getMessage.contains(Utils.makeShuffleKey(appId, shuffleId)))
}
val registeredShuffles = master.statusSystem.registeredAppAndShuffles.get(appId)
shuffleIds.forEach { shuffleId =>
assert(!registeredShuffles.contains(shuffleId))
}
assert(registeredShuffles.contains(remainingShuffleId))

response = webTarget.path("shuffles/unregister").request(MediaType.APPLICATION_JSON).post(
Entity.entity(request, MediaType.APPLICATION_JSON))
assert(HttpServletResponse.SC_OK == response.getStatus)
assert(response.readEntity(classOf[HandleResponse]).getSuccess)
assert(master.statusSystem.registeredAppAndShuffles.get(appId).contains(remainingShuffleId))
} finally {
master.statusSystem.registeredAppAndShuffles.remove(appId)
master.statusSystem.appHeartbeatTime.remove(appId)
}
}

test("unregister shuffles validates request") {
var response =
webTarget.path("shuffles/unregister").request(MediaType.APPLICATION_JSON).post(
Entity.entity("null", MediaType.APPLICATION_JSON))
assert(HttpServletResponse.SC_BAD_REQUEST == response.getStatus)
assert(response.readEntity(classOf[String]).contains("request is required"))

response =
webTarget.path("shuffles/unregister").request(MediaType.APPLICATION_JSON).post(
Entity.entity(
new UnregisterShufflesRequest().shuffleIds(Arrays.asList[Integer](1)),
MediaType.APPLICATION_JSON))
assert(HttpServletResponse.SC_BAD_REQUEST == response.getStatus)
assert(response.readEntity(classOf[String]).contains("appId"))

response = webTarget.path("shuffles/unregister").request(MediaType.APPLICATION_JSON).post(
Entity.entity(
new UnregisterShufflesRequest().appId(" ").shuffleIds(Arrays.asList[Integer](1)),
MediaType.APPLICATION_JSON))
assert(HttpServletResponse.SC_BAD_REQUEST == response.getStatus)
assert(response.readEntity(classOf[String]).contains("appId"))

response = webTarget.path("shuffles/unregister").request(MediaType.APPLICATION_JSON).post(
Entity.entity(
new UnregisterShufflesRequest().appId("app"),
MediaType.APPLICATION_JSON))
assert(HttpServletResponse.SC_BAD_REQUEST == response.getStatus)
assert(response.readEntity(classOf[String]).contains("nonempty"))

response = webTarget.path("shuffles/unregister").request(MediaType.APPLICATION_JSON).post(
Entity.entity(
new UnregisterShufflesRequest().appId("app").shuffleIds(Arrays.asList[Integer](-1)),
MediaType.APPLICATION_JSON))
assert(HttpServletResponse.SC_BAD_REQUEST == response.getStatus)
assert(response.readEntity(classOf[String]).contains("nonnegative"))
}

test("application resource") {
var response = webTarget.path("applications").request(MediaType.APPLICATION_JSON).get()
assert(HttpServletResponse.SC_OK == response.getStatus)
assert(response.readEntity(classOf[ApplicationsResponse]).getApplications.isEmpty)
assert(response.readEntity(classOf[ApplicationsHeartbeatResponse]).getApplications.isEmpty)

response = webTarget.path("applications/hostnames").request(MediaType.APPLICATION_JSON).get()
assert(HttpServletResponse.SC_OK == response.getStatus)
Expand Down
Loading
Loading