[ci] Move simulate_storage from _private/ to _common/ (#53735)

Used by Train and Tune.

Only `s3` type was used, so refactored.

---------

Signed-off-by: Edward Oakes <ed.nmi.oakes@gmail.com>

Author: edoakes

ci/lint/pydoclint-baseline.txt

@@ -111,16 +111,6 @@ python/ray/_private/node.py
     DOC107: Method `Node.kill_all_processes`: The option `--arg-type-hints-in-signature` is `True` but not all args in the signature have type hints
     DOC103: Method `Node.kill_all_processes`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [allow_graceful: ].
 --------------------
-python/ray/_private/parameter.py
-    DOC101: Method `RayParams.__init__`: Docstring contains fewer arguments than in function signature.
-    DOC107: Method `RayParams.__init__`: The option `--arg-type-hints-in-signature` is `True` but not all args in the signature have type hints
-    DOC103: Method `RayParams.__init__`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [_system_config: Optional[Dict[str, str]], autoscaling_config: Optional[str], cluster_id: Optional[str], dashboard_agent_listen_port: Optional[int], dashboard_host: Optional[str], dashboard_port: Optional[bool], driver_mode: , enable_object_reconstruction: Optional[bool], env_vars: Optional[Dict[str, str]], external_addresses: Optional[List[str]], gcs_address: Optional[str], gcs_server_port: Optional[int], huge_pages: Optional[bool], include_dashboard: Optional[bool], include_log_monitor: Optional[str], labels: Optional[Dict[str, str]], max_worker_port: Optional[int], memory: Optional[float], metrics_agent_port: Optional[int], metrics_export_port: Optional[int], min_worker_port: Optional[int], no_monitor: Optional[bool], node_id: Optional[str], node_ip_address: Optional[str], node_manager_port: int, node_name: Optional[str], num_cpus: Optional[int], num_gpus: Optional[int], num_redis_shards: Optional[int], object_manager_port: Optional[int], object_ref_seed: Optional[int], object_spilling_directory: Optional[str], object_store_memory: Optional[float], plasma_directory: Optional[str], plasma_store_socket_name: Optional[str], ray_client_server_port: Optional[int], ray_debugger_external: bool, raylet_ip_address: Optional[str], raylet_socket_name: Optional[str], redirect_output: Optional[bool], redis_address: Optional[str], redis_max_clients: Optional[int], redis_password: Optional[str], redis_port: Optional[int], redis_shard_ports: Optional[List[int]], redis_username: Optional[str], resource_isolation_config: Optional[ResourceIsolationConfig], resources: Optional[Dict[str, float]], runtime_env_agent_port: Optional[int], runtime_env_dir_name: Optional[str], session_name: Optional[str], setup_worker_path: Optional[str], storage: Optional[str], temp_dir: Optional[str], tracing_startup_hook: , webui: Optional[str], worker_path: Optional[str], worker_port_list: Optional[List[int]]].
-    DOC106: Method `RayParams.update`: The option `--arg-type-hints-in-signature` is `True` but there are no argument type hints in the signature
-    DOC103: Method `RayParams.update`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [**kwargs: ]. Arguments in the docstring but not in the function signature: [kwargs: ].
-    DOC106: Method `RayParams.update_if_absent`: The option `--arg-type-hints-in-signature` is `True` but there are no argument type hints in the signature
-    DOC103: Method `RayParams.update_if_absent`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [**kwargs: ]. Arguments in the docstring but not in the function signature: [kwargs: ].
-    DOC202: Method `RayParams.update_pre_selected_port` has a return section in docstring, but there are no return statements or annotations
---------------------
 python/ray/_private/profiling.py
     DOC106: Function `profile`: The option `--arg-type-hints-in-signature` is `True` but there are no argument type hints in the signature
     DOC107: Function `profile`: The option `--arg-type-hints-in-signature` is `True` but not all args in the signature have type hints
@@ -285,8 +275,6 @@ python/ray/_private/test_utils.py
     DOC201: Method `BatchQueue.get_batch` does not have a return section in docstring
     DOC101: Function `monitor_memory_usage`: Docstring contains fewer arguments than in function signature.
     DOC103: Function `monitor_memory_usage`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [print_interval_s: int, record_interval_s: int]. Arguments in the docstring but not in the function signature: [interval_s: ].
-    DOC402: Function `simulate_storage` has "yield" statements, but the docstring does not have a "Yields" section
-    DOC404: Function `simulate_storage` yield type(s) in docstring not consistent with the return annotation. Return annotation exists, but docstring "yields" section does not exist or has 0 type(s).
 --------------------
 python/ray/_private/usage/usage_lib.py
     DOC201: Function `record_extra_usage_tag` does not have a return section in docstring
@@ -2949,56 +2937,3 @@ python/ray/widgets/util.py
     DOC103: Function `_has_missing`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [*deps: Iterable[Union[str, Optional[str]]]]. Arguments in the docstring but not in the function signature: [deps: ].
     DOC103: Function `repr_with_fallback`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [*notebook_deps: Iterable[Union[str, Optional[str]]]]. Arguments in the docstring but not in the function signature: [notebook_deps: ].
 --------------------
-python/ray/workflow/api.py
-    DOC101: Function `run`: Docstring contains fewer arguments than in function signature.
-    DOC103: Function `run`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [**kwargs: , *args: , dag: DAGNode].
-    DOC101: Function `run_async`: Docstring contains fewer arguments than in function signature.
-    DOC103: Function `run_async`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [**kwargs: , *args: , dag: DAGNode].
-    DOC201: Function `continuation` does not have a return section in docstring
-    DOC101: Method `options.__init__`: Docstring contains fewer arguments than in function signature.
-    DOC103: Method `options.__init__`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [**workflow_options: Dict[str, Any]].
---------------------
-python/ray/workflow/serialization.py
-    DOC101: Method `Manager.save_objectref`: Docstring contains fewer arguments than in function signature.
-    DOC103: Method `Manager.save_objectref`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [workflow_id: 'str'].
-    DOC107: Function `dump_to_storage`: The option `--arg-type-hints-in-signature` is `True` but not all args in the signature have type hints
---------------------
-python/ray/workflow/serialization_context.py
-    DOC402: Function `workflow_args_serialization_context` has "yield" statements, but the docstring does not have a "Yields" section
-    DOC404: Function `workflow_args_serialization_context` yield type(s) in docstring not consistent with the return annotation. Return annotation exists, but docstring "yields" section does not exist or has 0 type(s).
-    DOC402: Function `workflow_args_resolving_context` has "yield" statements, but the docstring does not have a "Yields" section
-    DOC404: Function `workflow_args_resolving_context` yield type(s) in docstring not consistent with the return annotation. Return annotation exists, but docstring "yields" section does not exist or has 0 type(s).
---------------------
-python/ray/workflow/storage/filesystem.py
-    DOC107: Function `_open_atomic`: The option `--arg-type-hints-in-signature` is `True` but not all args in the signature have type hints
-    DOC202: Function `_open_atomic` has a return section in docstring, but there are no return statements or annotations
-    DOC402: Function `_open_atomic` has "yield" statements, but the docstring does not have a "Yields" section
-    DOC404: Function `_open_atomic` yield type(s) in docstring not consistent with the return annotation. Return annotation exists, but docstring "yields" section does not exist or has 0 type(s).
---------------------
-python/ray/workflow/task_executor.py
-    DOC104: Function `_workflow_task_executor`: Arguments are the same in the docstring and the function signature, but are in a different order.
-    DOC105: Function `_workflow_task_executor`: Argument names match, but type hints in these args do not match: func, context, task_id, baked_inputs, runtime_options
-    DOC101: Method `_BakedWorkflowInputs.resolve`: Docstring contains fewer arguments than in function signature.
-    DOC103: Method `_BakedWorkflowInputs.resolve`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [store: workflow_storage.WorkflowStorage].
---------------------
-python/ray/workflow/workflow_context.py
-    DOC106: Function `workflow_task_context`: The option `--arg-type-hints-in-signature` is `True` but there are no argument type hints in the signature
-    DOC107: Function `workflow_task_context`: The option `--arg-type-hints-in-signature` is `True` but not all args in the signature have type hints
-    DOC402: Function `workflow_task_context` has "yield" statements, but the docstring does not have a "Yields" section
-    DOC404: Function `workflow_task_context` yield type(s) in docstring not consistent with the return annotation. Return annotation exists, but docstring "yields" section does not exist or has 0 type(s).
-    DOC106: Function `workflow_logging_context`: The option `--arg-type-hints-in-signature` is `True` but there are no argument type hints in the signature
-    DOC107: Function `workflow_logging_context`: The option `--arg-type-hints-in-signature` is `True` but not all args in the signature have type hints
-    DOC402: Function `workflow_logging_context` has "yield" statements, but the docstring does not have a "Yields" section
-    DOC404: Function `workflow_logging_context` yield type(s) in docstring not consistent with the return annotation. Return annotation exists, but docstring "yields" section does not exist or has 0 type(s).
---------------------
-python/ray/workflow/workflow_state_from_dag.py
-    DOC201: Function `workflow_state_from_dag` does not have a return section in docstring
---------------------
-python/ray/workflow/workflow_storage.py
-    DOC201: Method `WorkflowIndexingStorage.list_workflow` does not have a return section in docstring
-    DOC201: Method `WorkflowStorage.load_actor_class_body` does not have a return section in docstring
-    DOC101: Method `WorkflowStorage.load_task_metadata`: Docstring contains fewer arguments than in function signature.
-    DOC103: Method `WorkflowStorage.load_task_metadata`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [task_id: TaskID].
-    DOC201: Method `WorkflowStorage.list_workflow` does not have a return section in docstring
-    DOC201: Method `WorkflowStorage._put` does not have a return section in docstring
---------------------

python/ray/_common/test_utils.py

@@ -6,10 +6,13 @@
 """
 
 import asyncio
+from contextlib import contextmanager
 import inspect
+import os
 import time
 import traceback
-from typing import Any, Callable
+from typing import Any, Callable, Iterator
+import uuid
 
 import ray
 import ray._private.utils
@@ -143,3 +146,34 @@ async def async_wait_for_condition(
     if last_ex is not None:
         message += f" Last exception: {last_ex}"
     raise RuntimeError(message)
+
+
+@contextmanager
+def simulate_s3_bucket(
+    port: int = 5002,
+    region: str = "us-west-2",
+) -> Iterator[str]:
+    """Context manager that simulates an S3 bucket and yields the URI.
+
+    Args:
+        port: The port of the localhost endpoint where S3 is being served.
+        region: The S3 region.
+
+    Yields:
+        str: URI for the simulated S3 bucket.
+    """
+    from moto.server import ThreadedMotoServer
+
+    old_env = os.environ
+    os.environ["AWS_ACCESS_KEY_ID"] = "testing"
+    os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
+    os.environ["AWS_SECURITY_TOKEN"] = "testing"
+    os.environ["AWS_SESSION_TOKEN"] = "testing"
+
+    s3_server = f"http://localhost:{port}"
+    server = ThreadedMotoServer(port=port)
+    server.start()
+    url = f"s3://{uuid.uuid4().hex}?region={region}&endpoint_override={s3_server}"
+    yield url
+    server.stop()
+    os.environ = old_env

python/ray/_private/test_utils.py

@@ -1774,50 +1774,6 @@ def no_resource_leaks_excluding_node_resources():
     return cluster_resources == available_resources
 
 
-@contextmanager
-def simulate_storage(
-    storage_type: str,
-    root: Optional[str] = None,
-    port: int = 5002,
-    region: str = "us-west-2",
-):
-    """Context that simulates a given storage type and yields the URI.
-
-    Args:
-        storage_type: The storage type to simiulate ("fs" or "s3")
-        root: Root directory of the URI to return (e.g., s3 bucket name)
-        port: The port of the localhost endpoint where s3 is being served (s3 only)
-        region: The s3 region (s3 only)
-    """
-    if storage_type == "fs":
-        if root is None:
-            with tempfile.TemporaryDirectory() as d:
-                yield "file://" + d
-        else:
-            yield "file://" + root
-    elif storage_type == "s3":
-        from moto.server import ThreadedMotoServer
-
-        old_env = os.environ
-        os.environ["AWS_ACCESS_KEY_ID"] = "testing"
-        os.environ["AWS_SECRET_ACCESS_KEY"] = "testing"
-        os.environ["AWS_SECURITY_TOKEN"] = "testing"
-        os.environ["AWS_SESSION_TOKEN"] = "testing"
-
-        root = root or uuid.uuid4().hex
-        s3_server = f"http://localhost:{port}"
-        server = ThreadedMotoServer(port=port)
-        server.start()
-        url = f"s3://{root}?region={region}&endpoint_override={s3_server}"
-        yield url
-        server.stop()
-
-        os.environ = old_env
-
-    else:
-        raise NotImplementedError(f"Unknown storage type: {storage_type}")
-
-
 def job_hook(**kwargs):
     """Function called by reflection by test_cli_integration."""
     cmd = " ".join(kwargs["entrypoint"])

python/ray/train/tests/conftest.py

@@ -4,7 +4,7 @@
 import pytest
 
 import ray
-from ray._private.test_utils import simulate_storage
+from ray._common.test_utils import simulate_s3_bucket
 from ray.cluster_utils import Cluster
 
 # Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
@@ -127,7 +127,7 @@ def mock_s3_bucket_uri():
 
     port = 5002
     region = "us-west-2"
-    with simulate_storage("s3", port=port, region=region) as s3_uri:
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
         s3 = boto3.client(
             "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
         )

python/ray/train/tests/test_new_persistence.py

@@ -14,7 +14,7 @@
 
 import ray
 from ray import train, tune
-from ray._private.test_utils import simulate_storage
+from ray._common.test_utils import simulate_s3_bucket
 from ray.air._internal.uri_utils import URI
 from ray.air.constants import EXPR_RESULT_FILE
 from ray.train._checkpoint import Checkpoint
@@ -40,7 +40,7 @@ class TestConstants:
 def mock_s3_bucket_uri():
     port = 5002
     region = "us-west-2"
-    with simulate_storage("s3", port=port, region=region) as s3_uri:
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
         import boto3
 
         s3 = boto3.client(

python/ray/train/v2/tests/test_persistence.py

@@ -14,7 +14,7 @@
 
 import ray
 import ray.train
-from ray._private.test_utils import simulate_storage
+from ray._common.test_utils import simulate_s3_bucket
 from ray.air._internal.uri_utils import URI
 from ray.train import (
     Checkpoint,
@@ -41,7 +41,7 @@ class TestConstants:
 def mock_s3_bucket_uri():
     port = 5002
     region = "us-west-2"
-    with simulate_storage("s3", port=port, region=region) as s3_uri:
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
         import boto3
 
         s3 = boto3.client(

python/ray/tune/tests/conftest.py

@@ -3,7 +3,7 @@
 import boto3
 import pytest
 
-from ray._private.test_utils import simulate_storage
+from ray._common.test_utils import simulate_s3_bucket
 from ray.air._internal.uri_utils import URI
 
 # Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
@@ -17,7 +17,7 @@
 def mock_s3_bucket_uri():
     port = 5002
     region = "us-west-2"
-    with simulate_storage("s3", port=port, region=region) as s3_uri:
+    with simulate_s3_bucket(port=port, region=region) as s3_uri:
         s3 = boto3.client(
             "s3", region_name=region, endpoint_url=f"http://localhost:{port}"
         )