FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

feat: Add cloud_function_cpus option to remote_function (#2475) · google/bigframes@4caf74c · GitHub

/ bigframes Public

Commit 4caf74c

Browse files
TrevorBergeron
authored
feat: Add cloud_function_cpus option to remote_function (#2475)
1 parent cb00daa commit 4caf74c

5 files changed

Lines changed: 120 additions & 17 deletions

File tree

‎bigframes/functions/_function_client.py‎

Lines changed: 67 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@
6565
# BQ managed functions (@udf) currently only support Python 3.11.
6666
_MANAGED_FUNC_PYTHON_VERSION = "python-3.11"
6767

68+
_DEFAULT_FUNCTION_MEMORY_MIB = 1024
69+
6870

6971
class FunctionClient:
7072
# Wait time (in seconds) for an IAM binding to take effect after creation.
@@ -402,8 +404,12 @@ def create_cloud_function(
402404
is_row_processor=False,
403405
vpc_connector=None,
404406
vpc_connector_egress_settings="private-ranges-only",
405-
memory_mib=1024,
407+
memory_mib=None,
408+
cpus=None,
406409
ingress_settings="internal-only",
410+
workers=None,
411+
threads=None,
412+
concurrency=None,
407413
):
408414
"""Create a cloud function from the given user defined function."""
409415

@@ -486,6 +492,8 @@ def create_cloud_function(
486492
function.service_config = functions_v2.ServiceConfig()
487493
if memory_mib is not None:
488494
function.service_config.available_memory = f"{memory_mib}Mi"
495+
if cpus is not None:
496+
function.service_config.available_cpu = str(cpus)
489497
if timeout_seconds is not None:
490498
if timeout_seconds > 1200:
491499
raise bf_formatting.create_exception_with_feedback_link(
@@ -517,6 +525,20 @@ def create_cloud_function(
517525
function.service_config.service_account_email = (
518526
self._cloud_function_service_account
519527
)
528+
if concurrency:
529+
function.service_config.max_instance_request_concurrency = concurrency
530+
531+
# Functions framework use environment variables to pass config to gunicorn
532+
# See https://github.com/GoogleCloudPlatform/functions-framework-python/issues/241
533+
# Code: https://github.com/GoogleCloudPlatform/functions-framework-python/blob/v3.10.1/src/functions_framework/_http/gunicorn.py#L37-L43
534+
env_vars = {}
535+
if workers:
536+
env_vars["WORKERS"] = str(workers)
537+
if threads:
538+
env_vars["THREADS"] = str(threads)
539+
if env_vars:
540+
function.service_config.environment_variables = env_vars
541+
520542
if ingress_settings not in _INGRESS_SETTINGS_MAP:
521543
raise bf_formatting.create_exception_with_feedback_link(
522544
ValueError,
@@ -581,6 +603,7 @@ def provision_bq_remote_function(
581603
cloud_function_vpc_connector,
582604
cloud_function_vpc_connector_egress_settings,
583605
cloud_function_memory_mib,
606+
cloud_function_cpus,
584607
cloud_function_ingress_settings,
585608
bq_metadata,
586609
):
@@ -616,6 +639,21 @@ def provision_bq_remote_function(
616639
)
617640
cf_endpoint = self.get_cloud_function_endpoint(cloud_function_name)
618641

642+
if cloud_function_memory_mib is None:
643+
cloud_function_memory_mib = _DEFAULT_FUNCTION_MEMORY_MIB
644+
645+
# assumption is most bigframes functions are cpu bound, single-threaded and many won't release GIL
646+
# therefore, want to allocate a worker for each cpu, and allow a concurrent request per worker
647+
expected_milli_cpus = (
648+
int(cloud_function_cpus * 1000)
649+
if (cloud_function_cpus is not None)
650+
else _infer_milli_cpus_from_memory(cloud_function_memory_mib)
651+
)
652+
workers = -(expected_milli_cpus // -1000) # ceil(cpus) without invoking floats
653+
threads = 4 # (per worker)
654+
# max concurrency==1 for vcpus < 1 hard limit from cloud run
655+
concurrency = (workers * threads) if (expected_milli_cpus >= 1000) else 1
656+
619657
# Create the cloud function if it does not exist
620658
if not cf_endpoint:
621659
cf_endpoint = self.create_cloud_function(
@@ -630,7 +668,11 @@ def provision_bq_remote_function(
630668
vpc_connector=cloud_function_vpc_connector,
631669
vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
632670
memory_mib=cloud_function_memory_mib,
671+
cpus=cloud_function_cpus,
633672
ingress_settings=cloud_function_ingress_settings,
673+
workers=workers,
674+
threads=threads,
675+
concurrency=concurrency,
634676
)
635677
else:
636678
logger.info(f"Cloud function {cloud_function_name} already exists.")
@@ -696,3 +738,27 @@ def get_remote_function_specs(self, remote_function_name):
696738
# Note: list_routines doesn't make an API request until we iterate on the response object.
697739
pass
698740
return (http_endpoint, bq_connection)
741+
742+
743+
def _infer_milli_cpus_from_memory(memory_mib: int) -> int:
744+
# observed values, not formally documented by cloud run functions
745+
if memory_mib < 128:
746+
raise ValueError("Cloud run supports at minimum 128MiB per instance")
747+
elif memory_mib == 128:
748+
return 83
749+
elif memory_mib <= 256:
750+
return 167
751+
elif memory_mib <= 512:
752+
return 333
753+
elif memory_mib <= 1024:
754+
return 583
755+
elif memory_mib <= 2048:
756+
return 1000
757+
elif memory_mib <= 8192:
758+
return 2000
759+
elif memory_mib <= 16384:
760+
return 4000
761+
elif memory_mib <= 32768:
762+
return 8000
763+
else:
764+
raise ValueError("Cloud run supports at most 32768MiB per instance")

‎bigframes/functions/_function_session.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -248,7 +248,8 @@ def remote_function(
248248
cloud_function_vpc_connector_egress_settings: Optional[
249249
Literal["all", "private-ranges-only", "unspecified"]
250250
] = None,
251-
cloud_function_memory_mib: Optional[int] = 1024,
251+
cloud_function_memory_mib: Optional[int] = None,
252+
cloud_function_cpus: Optional[float] = None,
252253
cloud_function_ingress_settings: Literal[
253254
"all", "internal-only", "internal-and-gclb"
254255
] = "internal-only",
@@ -444,6 +445,10 @@ def remote_function(
444445
default memory of cloud functions be allocated, pass `None`. See
445446
for more details
446447
https://cloud.google.com/functions/docs/configuring/memory.
448+
cloud_function_cpus (float, Optional):
449+
The number of cpus to allocate for the cloud
450+
function (2nd gen) created.
451+
https://docs.cloud.google.com/run/docs/configuring/services/cpu.
447452
cloud_function_ingress_settings (str, Optional):
448453
Ingress settings controls dictating what traffic can reach the
449454
function. Options are: `all`, `internal-only`, or `internal-and-gclb`.
@@ -638,6 +643,7 @@ def wrapper(func):
638643
cloud_function_vpc_connector=cloud_function_vpc_connector,
639644
cloud_function_vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
640645
cloud_function_memory_mib=cloud_function_memory_mib,
646+
cloud_function_cpus=cloud_function_cpus,
641647
cloud_function_ingress_settings=cloud_function_ingress_settings,
642648
bq_metadata=bqrf_metadata,
643649
)

‎bigframes/pandas/__init__.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,8 @@ def remote_function(
8888
cloud_function_vpc_connector_egress_settings: Optional[
8989
Literal["all", "private-ranges-only", "unspecified"]
9090
] = None,
91-
cloud_function_memory_mib: Optional[int] = 1024,
91+
cloud_function_memory_mib: Optional[int] = None,
92+
cloud_function_cpus: Optional[float] = None,
9293
cloud_function_ingress_settings: Literal[
9394
"all", "internal-only", "internal-and-gclb"
9495
] = "internal-only",
@@ -112,6 +113,7 @@ def remote_function(
112113
cloud_function_vpc_connector=cloud_function_vpc_connector,
113114
cloud_function_vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
114115
cloud_function_memory_mib=cloud_function_memory_mib,
116+
cloud_function_cpus=cloud_function_cpus,
115117
cloud_function_ingress_settings=cloud_function_ingress_settings,
116118
cloud_build_service_account=cloud_build_service_account,
117119
)

‎bigframes/session/__init__.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1536,7 +1536,8 @@ def remote_function(
15361536
cloud_function_vpc_connector_egress_settings: Optional[
15371537
Literal["all", "private-ranges-only", "unspecified"]
15381538
] = None,
1539-
cloud_function_memory_mib: Optional[int] = 1024,
1539+
cloud_function_memory_mib: Optional[int] = None,
1540+
cloud_function_cpus: Optional[float] = None,
15401541
cloud_function_ingress_settings: Literal[
15411542
"all", "internal-only", "internal-and-gclb"
15421543
] = "internal-only",
@@ -1717,6 +1718,10 @@ def remote_function(
17171718
default memory of cloud functions be allocated, pass `None`. See
17181719
for more details
17191720
https://cloud.google.com/functions/docs/configuring/memory.
1721+
cloud_function_cpus (float, Optional):
1722+
The number of cpus to allocate for the cloud
1723+
function (2nd gen) created.
1724+
https://docs.cloud.google.com/run/docs/configuring/services/cpu.
17201725
cloud_function_ingress_settings (str, Optional):
17211726
Ingress settings controls dictating what traffic can reach the
17221727
function. Options are: `all`, `internal-only`, or `internal-and-gclb`.
@@ -1767,6 +1772,7 @@ def remote_function(
17671772
cloud_function_vpc_connector=cloud_function_vpc_connector,
17681773
cloud_function_vpc_connector_egress_settings=cloud_function_vpc_connector_egress_settings,
17691774
cloud_function_memory_mib=cloud_function_memory_mib,
1775+
cloud_function_cpus=cloud_function_cpus,
17701776
cloud_function_ingress_settings=cloud_function_ingress_settings,
17711777
cloud_build_service_account=cloud_build_service_account,
17721778
)

‎tests/system/large/functions/test_remote_function.py‎

Lines changed: 36 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2089,19 +2089,40 @@ def foo_list(x: pandas.Series, y0: float, y1, y2) -> list[str]:
20892089

20902090

20912091
@pytest.mark.parametrize(
2092-
("memory_mib_args", "expected_memory"),
2092+
(
2093+
"memory_mib_args",
2094+
"expected_memory",
2095+
"expected_cpus",
2096+
),
20932097
[
2094-
pytest.param({}, "1024Mi", id="no-set"),
2095-
pytest.param({"cloud_function_memory_mib": None}, "256M", id="set-None"),
2096-
pytest.param({"cloud_function_memory_mib": 128}, "128Mi", id="set-128"),
2097-
pytest.param({"cloud_function_memory_mib": 1024}, "1024Mi", id="set-1024"),
2098-
pytest.param({"cloud_function_memory_mib": 4096}, "4096Mi", id="set-4096"),
2099-
pytest.param({"cloud_function_memory_mib": 32768}, "32768Mi", id="set-32768"),
2098+
pytest.param({}, "1024Mi", None, id="no-set"),
2099+
pytest.param(
2100+
{"cloud_function_memory_mib": None}, "1024Mi", None, id="set-None"
2101+
),
2102+
pytest.param({"cloud_function_memory_mib": 128}, "128Mi", None, id="set-128"),
2103+
pytest.param(
2104+
{"cloud_function_memory_mib": 512, "cloud_function_cpus": 0.6},
2105+
"512Mi",
2106+
"0.6",
2107+
id="set-512",
2108+
),
2109+
pytest.param(
2110+
{"cloud_function_memory_mib": 1024}, "1024Mi", None, id="set-1024"
2111+
),
2112+
pytest.param(
2113+
{"cloud_function_memory_mib": 4096, "cloud_function_cpus": 4},
2114+
"4096Mi",
2115+
"4",
2116+
id="set-4096",
2117+
),
2118+
pytest.param(
2119+
{"cloud_function_memory_mib": 32768}, "32768Mi", None, id="set-32768"
2120+
),
21002121
],
21012122
)
21022123
@pytest.mark.flaky(retries=2, delay=120)
21032124
def test_remote_function_gcf_memory(
2104-
session, scalars_dfs, memory_mib_args, expected_memory
2125+
session, scalars_dfs, memory_mib_args, expected_memory, expected_cpus
21052126
):
21062127
try:
21072128

@@ -2117,6 +2138,12 @@ def square(x: int) -> int:
21172138
name=square_remote.bigframes_cloud_function
21182139
)
21192140
assert gcf.service_config.available_memory == expected_memory
2141+
if expected_cpus is not None:
2142+
assert gcf.service_config.available_cpu == expected_cpus
2143+
if float(gcf.service_config.available_cpu) >= 1.0:
2144+
assert gcf.service_config.max_instance_request_concurrency >= float(
2145+
gcf.service_config.available_cpu
2146+
)
21202147

21212148
scalars_df, scalars_pandas_df = scalars_dfs
21222149

@@ -2138,12 +2165,8 @@ def square(x: int) -> int:
21382165
pytest.param(32769, id="set-32769-too-high"),
21392166
],
21402167
)
2141-
@pytest.mark.flaky(retries=2, delay=120)
21422168
def test_remote_function_gcf_memory_unsupported(session, memory_mib):
2143-
with pytest.raises(
2144-
google.api_core.exceptions.InvalidArgument,
2145-
match="Invalid value specified for container memory",
2146-
):
2169+
with pytest.raises(ValueError, match="Cloud run supports"):
21472170

21482171
@session.remote_function(
21492172
reuse=False,

0 commit comments

Comments
 (0)
Back | FazBrowse Home | New Git URL