zhengruifeng commented on code in PR #58264: URL: https://github.com/apache/spark/pull/58264#discussion_r3859213317
########## sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/PythonWorkerEnvironment.scala: ########## @@ -0,0 +1,219 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connect.service + +import java.nio.charset.StandardCharsets + +import org.apache.spark.{SparkEnv, SparkException} +import org.apache.spark.sql.RuntimeConfig +import org.apache.spark.sql.connect.config.Connect +import org.apache.spark.sql.internal.SQLConf + +/** + * The environment variables that Python worker processes launched for a session's Python + * functions should inherit. + * + * The environment is carried by session configurations under a reserved prefix, one configuration + * per variable: `spark.pythonWorkerEnv.FOO=bar` makes `FOO` visible as `bar` in `os.environ` + * inside a Python UDF. The configurations are the authoritative session state -- no second copy + * of the environment is maintained as session state -- so the environment follows the session + * wherever ordinary session configurations follow it, including into a session created by + * `cloneSession`. A request's snapshot is also held in the plan cache keys of the plans it + * caches, since a cached plan is only reusable by a request carrying the same environment. + * + * Names are preserved case-sensitively by Spark. On a case-sensitive operating system `FOO` and + * `foo` are therefore distinct variables; Windows process environments are case-insensitive, so + * what a worker observes there is the platform's business rather than Spark's. + * + * A request reads the environment once and uses that one snapshot for everything it does, because + * the configurations can change underneath it: another request may set them while this one is + * still planning. Re-reading would let a plan be built with one environment and cached under + * another. + */ +private[connect] object PythonWorkerEnvironment { + + /** Prefix of the session configurations that carry the environment. */ + val confPrefix: String = "spark.pythonWorkerEnv." + + /** + * Environment variable names accepted under [[confPrefix]]. + * + * This is deliberately stricter than the operating system requires. A POSIX environment permits + * any byte except `=` and NUL in a name, and container platforms accept their own broader sets, + * but a name outside this pattern cannot be referenced portably from a shell, so accepting one + * would let a session install a variable that some consumers can never read. It is a + * portability policy, not a description of what a process environment can hold. + */ + val namePattern: String = "^[A-Za-z_][A-Za-z0-9_]*$" Review Comment: Spark-owned worker controls need a collision policy before this prefix accepts arbitrary names. For example, when `reuseWorker` is false the runner does not overwrite `SPARK_REUSE_WORKER`, so a session-provided value survives and the Python daemon enables reuse. Please reject Spark-owned names here (or explicitly clear every conditional control before launch) and add collision tests beyond the unconditional `PYTHONUNBUFFERED` case. ########## sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala: ########## @@ -87,6 +87,15 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.Utils +/** + * Translates a Spark Connect request into Catalyst. + * + * An instance is request-scoped: construct one per request and discard it. Some state is derived + * once and reused for the whole request -- notably the Python worker environment, which must be a + * single snapshot so that a plan cannot be built with one environment and cached under another. + * Reusing an instance across requests would pin that state to whatever the first request Review Comment: Capture the environment per top-level transform instead of making this reusable DeveloperApi object one-shot by convention. `SparkConnectPlanner` remains publicly constructible and `transformRelation` can be called repeatedly, so a caller that updates the session environment between transforms silently reuses the first snapshot and cache key. Please thread a request snapshot through recursive translation, or structurally enforce a one-request planner lifetime. ########## sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/PythonWorkerEnvironment.scala: ########## @@ -0,0 +1,219 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connect.service + +import java.nio.charset.StandardCharsets + +import org.apache.spark.{SparkEnv, SparkException} +import org.apache.spark.sql.RuntimeConfig +import org.apache.spark.sql.connect.config.Connect +import org.apache.spark.sql.internal.SQLConf + +/** + * The environment variables that Python worker processes launched for a session's Python + * functions should inherit. + * + * The environment is carried by session configurations under a reserved prefix, one configuration + * per variable: `spark.pythonWorkerEnv.FOO=bar` makes `FOO` visible as `bar` in `os.environ` + * inside a Python UDF. The configurations are the authoritative session state -- no second copy + * of the environment is maintained as session state -- so the environment follows the session + * wherever ordinary session configurations follow it, including into a session created by + * `cloneSession`. A request's snapshot is also held in the plan cache keys of the plans it + * caches, since a cached plan is only reusable by a request carrying the same environment. + * + * Names are preserved case-sensitively by Spark. On a case-sensitive operating system `FOO` and + * `foo` are therefore distinct variables; Windows process environments are case-insensitive, so + * what a worker observes there is the platform's business rather than Spark's. + * + * A request reads the environment once and uses that one snapshot for everything it does, because + * the configurations can change underneath it: another request may set them while this one is + * still planning. Re-reading would let a plan be built with one environment and cached under + * another. + */ +private[connect] object PythonWorkerEnvironment { + + /** Prefix of the session configurations that carry the environment. */ + val confPrefix: String = "spark.pythonWorkerEnv." + + /** + * Environment variable names accepted under [[confPrefix]]. + * + * This is deliberately stricter than the operating system requires. A POSIX environment permits + * any byte except `=` and NUL in a name, and container platforms accept their own broader sets, + * but a name outside this pattern cannot be referenced portably from a shell, so accepting one + * would let a session install a variable that some consumers can never read. It is a + * portability policy, not a description of what a process environment can hold. + */ + val namePattern: String = "^[A-Za-z_][A-Za-z0-9_]*$" + + private val compiledNamePattern = namePattern.r + + // A rejected name can be arbitrarily long, so messages carry a bounded prefix of it rather than + // the whole name. + private val maxNameCharsInMessage = 32 + + /** + * The environment carried by `conf`, without validation. + * + * Callers take one snapshot per request and pass it around. Validation is separate so that the + * plan cache can tell two environments apart without rejecting an invalid one: an invalid entry + * has to fail the queries that would install it in a worker, not every query in the session. + */ + def read(conf: SQLConf): Map[String, String] = extract(conf.getAllConfs) + + /** The environment carried by the configurations in `allConfs`. */ + private def extract(allConfs: Map[String, String]): Map[String, String] = { + allConfs.iterator + .filter { case (key, _) => key.startsWith(confPrefix) } + .map { case (key, value) => key.substring(confPrefix.length) -> value } + .toMap + } + + /** + * Rejects a malformed or oversized environment. + * + * This runs when a Python function is built, which is the one point that every way of writing a + * configuration reaches: the Spark Connect config RPC, SQL `SET`, and the application-level + * configurations merged into a new session all arrive here. [[validateConfigChange]] rejects a + * write through the config RPC earlier and more helpfully, but it cannot see the other two, so + * this is the check that makes an invalid environment unable to reach a worker at all. + * + * A message may name a variable but never carries its value, so a rejection cannot copy a value + * into a log or a stack trace. Note that the name is chosen by the user, so a name is only as + * safe as what the user put in it. + * + * @throws SparkException + * if a name is malformed or too long, a value cannot be carried by a process environment, or + * the collection exceeds a limit. + */ + def validate(variables: Map[String, String]): Unit = { + val conf = SparkEnv.get.conf + val maxCount = conf.get(Connect.CONNECT_PYTHON_WORKER_ENV_MAX_VARIABLES) + val maxNameLength = conf.get(Connect.CONNECT_PYTHON_WORKER_ENV_MAX_NAME_LENGTH) + val maxTotalSizeBytes = conf.get(Connect.CONNECT_PYTHON_WORKER_ENV_MAX_TOTAL_SIZE_BYTES) + + if (variables.size > maxCount) { + throw new SparkException( + errorClass = "INVALID_SPARK_CONFIG.PYTHON_WORKER_ENV_TOO_MANY_VARIABLES", + messageParameters = Map( + "count" -> variables.size.toString, + "prefix" -> confPrefix, + "maxCount" -> maxCount.toString), + cause = null) + } + + var totalSizeBytes = 0L + variables.foreach { case (name, value) => + // `matches` requires the whole name to match. Searching for the pattern instead would accept + // a name with a trailing newline, because `$` also matches before a terminating line break. + if (name.length > maxNameLength || !compiledNamePattern.matches(name)) { + throw new SparkException( + errorClass = "INVALID_SPARK_CONFIG.INVALID_PYTHON_WORKER_ENV_VAR_NAME", + messageParameters = Map( + "name" -> describeName(name), + "prefix" -> confPrefix, + "pattern" -> namePattern, + "maxLength" -> maxNameLength.toString), + cause = null) + } + // A process environment cannot carry NUL. Rejecting it here rather than letting the worker + // launch fail matters for more than the error message: the launch failure is an + // `IllegalArgumentException` from the JDK whose own message embeds the offending value. + if (value.indexOf(0) >= 0) { + throw new SparkException( + errorClass = "INVALID_SPARK_CONFIG.INVALID_PYTHON_WORKER_ENV_VAR_VALUE", + messageParameters = Map("name" -> describeName(name), "prefix" -> confPrefix), + cause = null) + } + totalSizeBytes += utf8Length(name) + utf8Length(value) + } + + if (totalSizeBytes > maxTotalSizeBytes) { + throw new SparkException( + errorClass = "INVALID_SPARK_CONFIG.PYTHON_WORKER_ENV_TOO_LARGE", + messageParameters = Map( + "prefix" -> confPrefix, + "size" -> totalSizeBytes.toString, + "maxSize" -> maxTotalSizeBytes.toString), + cause = null) + } + } + + /** + * Rejects a configuration write that would leave the session with an invalid environment. + * + * A no-op for a key outside [[confPrefix]]. For a key under it, the environment that the write + * would produce is validated before the write happens, so an invalid environment never enters + * the session at all and the failure points at the call that caused it. + * + * This covers the Spark Connect config RPC, which is both how a client sets a configuration + * explicitly and how `SparkSession.builder.config` applies one. It does not cover SQL `SET` or + * the application-level configurations merged into a new session: both reach the session + * configurations without passing through the RPC, which is why [[validate]] at build time stays + * as the check that no invalid environment can reach a worker. + * + * Removing a variable is deliberately not validated. A removal can only shrink the environment, + * and it is how a session recovers from an environment that one of those unchecked paths left + * invalid; validating a removal would leave such a session with no way back. + * + * @throws SparkException + * if the environment the write would produce is malformed or oversized. + */ + def validateConfigChange(conf: RuntimeConfig, key: String, value: Option[String]): Unit = { + if (key.startsWith(confPrefix)) { + // An absent value is rejected by `SQLConf` itself. Leave that failure where it is instead of + // reporting a missing value as an invalid environment. + value.foreach { newValue => + validate(extract(conf.getAll) + (key.substring(confPrefix.length) -> newValue)) Review Comment: A config RPC should reject any write that leaves the environment over its cluster limit, including under concurrent requests. This snapshots `getAll`, but `handleSet` performs `conf.set` under a separate synchronization; two writers can both validate against the old map and then jointly exceed `maxVariables` or `maxTotalSizeBytes`. Please make the full read/validate/set sequence atomic on SQLConf's settings lock and add a barrier-based concurrent-write test. ########## sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SparkConnectConfigHandler.scala: ########## @@ -85,6 +85,10 @@ class SparkConnectConfigHandler(responseObserver: StreamObserver[proto.ConfigRes operation.getPairsList.asScala.iterator.foreach { pair => val (key, value) = SparkConnectConfigHandler.toKeyValue(pair) try { + // Reject a write that would leave the session's Python worker environment invalid, before + // it is stored. Inside the try so that a `silent` request reports it as a warning, the way + // it reports any other rejected write. + PythonWorkerEnvironment.validateConfigChange(conf, key, value) Review Comment: Rejected environment values must stay out of the silent path too. This validation can throw into the catch below, whose warning interpolates the full `$value`; builder configs use silent mode, and the Python/JVM clients emit or log that warning. A NUL-containing secret is therefore disclosed despite `validate` omitting values. Please remove or redact the value and cover a silent NUL rejection. ########## sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/PythonWorkerEnvironment.scala: ########## @@ -0,0 +1,219 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.sql.connect.service + +import java.nio.charset.StandardCharsets + +import org.apache.spark.{SparkEnv, SparkException} +import org.apache.spark.sql.RuntimeConfig +import org.apache.spark.sql.connect.config.Connect +import org.apache.spark.sql.internal.SQLConf + +/** + * The environment variables that Python worker processes launched for a session's Python + * functions should inherit. + * + * The environment is carried by session configurations under a reserved prefix, one configuration + * per variable: `spark.pythonWorkerEnv.FOO=bar` makes `FOO` visible as `bar` in `os.environ` Review Comment: Please add `spark.pythonWorkerEnv.<NAME>` to the published configuration documentation. This private ScalaDoc is currently the only explanation of the user-facing entry point, so users cannot discover its limits, validation rules, SQL `SET` failure timing, or how it differs from application-scoped `spark.executorEnv.*`. ########## sql/connect/server/src/main/scala/org/apache/spark/sql/connect/service/SessionHolder.scala: ########## @@ -649,9 +660,24 @@ case class SessionHolder(userId: String, sessionId: String, session: SparkSessio } // For testing. Expose the plan cache for testing purposes. - private[service] def getPlanCache: Option[Cache[proto.Relation, LogicalPlan]] = planCache + private[service] def getPlanCache: Option[Cache[PlanCacheKey, LogicalPlan]] = planCache } +/** + * Key of an entry in a session's plan cache. + * + * @param relation + * the relation the cached plan was built from. + * @param pythonWorkerEnv + * the Python worker environment the plan was built with, as the request that built it observed + * it. Part of the key because the environment is baked into every Python function the plan Review Comment: Please change `Part of the key because ...` to `It is part of the key because ...`; the current `@param` text is a sentence fragment in the generated Scaladoc. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
