[incubator-celeborn] branch main updated: [CELEBORN-140][IMPROVEMENT] Quota yaml should support byte string (#1085)

[zhouky]

This is an automated email from the ASF dual-hosted git repository.

zhouky pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git


The following commit(s) were added to refs/heads/main by this push:
     new ed26356c [CELEBORN-140][IMPROVEMENT] Quota yaml should support byte 
string (#1085)
ed26356c is described below

commit ed26356cf7de7bd2a1c0fe5ad5b08083699846a4
Author: Angerszhuuuu <angers....@gmail.com>
AuthorDate: Wed Dec 14 18:24:19 2022 +0800

    [CELEBORN-140][IMPROVEMENT] Quota yaml should support byte string (#1085)
---
 .../org/apache/celeborn/common/quota/DefaultQuotaManager.scala    | 7 ++++++-
 common/src/test/resources/test-quota.yaml                         | 8 ++------
 .../apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala   | 5 +++--
 3 files changed, 11 insertions(+), 9 deletions(-)

diff --git 
a/common/src/main/scala/org/apache/celeborn/common/quota/DefaultQuotaManager.scala
 
b/common/src/main/scala/org/apache/celeborn/common/quota/DefaultQuotaManager.scala
index 300b3f41..fb327d33 100644
--- 
a/common/src/main/scala/org/apache/celeborn/common/quota/DefaultQuotaManager.scala
+++ 
b/common/src/main/scala/org/apache/celeborn/common/quota/DefaultQuotaManager.scala
@@ -49,7 +49,12 @@ class DefaultQuotaManager(conf: CelebornConf) extends 
QuotaManager(conf) {
             .asInstanceOf[java.util.HashMap[String, Object]]
             .asScala
             .foreach { case (key, value) =>
-              quota.update(userIdentifier, key, value.toString.toLong)
+              key match {
+                case "diskBytesWritten" | "hdfsBytesWritten" =>
+                  quota.update(userIdentifier, key, 
Utils.byteStringAsBytes(value.toString))
+                case _ =>
+                  quota.update(userIdentifier, key, value.toString.toLong)
+              }
             }
           userQuotas.put(userIdentifier, quota)
         }
diff --git a/common/src/test/resources/test-quota.yaml 
b/common/src/test/resources/test-quota.yaml
index 45d2cdc3..95935188 100644
--- a/common/src/test/resources/test-quota.yaml
+++ b/common/src/test/resources/test-quota.yaml
@@ -18,15 +18,11 @@
 -  tenantId: AAA
    name: Tom
    quota:
-     diskBytesWritten: 10000
+     diskBytesWritten: 100m
      diskFileCount: 200
-     hdfsBytesWritten: -1
-     hdfsFileCount: -1
 
 -  tenantId: BBB
    name: Jerry
    quota:
-     diskBytesWritten: -1
-     diskFileCount: -1
-     hdfsBytesWritten: 10000
+     hdfsBytesWritten: 200m
      hdfsFileCount: 200
diff --git 
a/common/src/test/scala/org/apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala
 
b/common/src/test/scala/org/apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala
index e7d43ef0..a1219ef6 100644
--- 
a/common/src/test/scala/org/apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala
+++ 
b/common/src/test/scala/org/apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala
@@ -21,6 +21,7 @@ import org.junit.Assert.assertEquals
 
 import org.apache.celeborn.common.CelebornConf
 import org.apache.celeborn.common.identity.UserIdentifier
+import org.apache.celeborn.common.util.Utils
 
 class DefaultQuotaManagerSuite extends BaseQuotaManagerSuite {
 
@@ -39,9 +40,9 @@ class DefaultQuotaManagerSuite extends BaseQuotaManagerSuite {
   test("test rss quota conf") {
     assertEquals(
       quotaManager.getQuota(UserIdentifier("AAA", "Tom")),
-      Quota(10000, 200, -1, -1))
+      Quota(Utils.byteStringAsBytes("100m"), 200, -1, -1))
     assertEquals(
       quotaManager.getQuota(UserIdentifier("BBB", "Jerry")),
-      Quota(-1, -1, 10000, 200))
+      Quota(-1, -1, Utils.byteStringAsBytes("200m"), 200))
   }
 }