chore: import upstream snapshot with attribution

python/ray/tests/runtime_env_container/test_job.py

@@ -0,0 +1,38 @@
+import argparse
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.job_submission import JobStatus, JobSubmissionClient
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+args = parser.parse_args()
+
+if args.use_image_uri_api:
+    runtime_env = {"image_uri": args.image}
+else:
+    runtime_env = {"container": {"image": args.image}}
+
+
+ray.init()
+client = JobSubmissionClient()
+job_id = client.submit_job(
+    entrypoint="cat file.txt",
+    runtime_env=runtime_env,
+)
+
+
+def check_job_succeeded():
+    assert client.get_job_status(job_id) == JobStatus.SUCCEEDED
+    return True
+
+
+wait_for_condition(check_job_succeeded)
+logs = client.get_job_logs(job_id)
+print("Job Logs:", logs)
+assert "helloworldalice" in logs

python/ray/tests/runtime_env_container/test_log_file_exists.py

@@ -0,0 +1,59 @@
+import argparse
+import re
+from pathlib import Path
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray.util.state import list_tasks
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+args = parser.parse_args()
+
+ray.init(num_cpus=1)
+
+session_dir = ray._private.worker.global_worker.node.address_info["session_dir"]
+session_path = Path(session_dir)
+log_dir_path = session_path / "logs"
+
+
+def task_finished():
+    tasks = list_tasks()
+    assert len(tasks) > 0
+    assert tasks[0].worker_id
+    assert tasks[0].worker_pid
+    assert tasks[0].state == "FINISHED"
+    return True
+
+
+if args.use_image_uri_api:
+    runtime_env = {"image_uri": args.image}
+else:
+    runtime_env = {"container": {"image": args.image}}
+
+
+# Run a basic workload.
+@ray.remote(runtime_env=runtime_env)
+def f():
+    for i in range(10):
+        print(f"test {i}")
+
+
+f.remote()
+wait_for_condition(task_finished)
+
+task_state = list_tasks()[0]
+worker_id = task_state.worker_id
+worker_pid = task_state.worker_pid
+print(f"Worker ID: {worker_id}")
+print(f"Worker PID: {worker_pid}")
+
+paths = [path.name for path in log_dir_path.iterdir()]
+assert f"python-core-worker-{worker_id}_{worker_pid}.log" in paths
+assert any(re.search(f"^worker-{worker_id}-.*-{worker_pid}.err$", p) for p in paths)
+assert any(re.search(f"^worker-{worker_id}-.*-{worker_pid}.out$", p) for p in paths)

python/ray/tests/runtime_env_container/test_put_get.py

@@ -0,0 +1,33 @@
+import argparse
+
+import numpy as np
+
+import ray
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+args = parser.parse_args()
+
+
+if args.use_image_uri_api:
+    runtime_env = {"image_uri": args.image}
+else:
+    runtime_env = {"container": {"image": args.image}}
+
+
+@ray.remote(runtime_env=runtime_env)
+def create_ref():
+    with open("file.txt") as f:
+        assert f.read().strip() == "helloworldalice"
+
+    ref = ray.put(np.zeros(100_000_000))
+    return ref
+
+
+wrapped_ref = create_ref.remote()
+assert (ray.get(ray.get(wrapped_ref)) == np.zeros(100_000_000)).all()

python/ray/tests/runtime_env_container/test_ray_env_vars.py

@@ -0,0 +1,33 @@
+import argparse
+import os
+
+import ray
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+args = parser.parse_args()
+
+
+if args.use_image_uri_api:
+    runtime_env = {"image_uri": args.image}
+else:
+    runtime_env = {"container": {"image": args.image}}
+
+
+@ray.remote(runtime_env=runtime_env)
+def f():
+    return os.environ.get("RAY_TEST_ABC")
+
+
+@ray.remote(runtime_env=runtime_env)
+def g():
+    return os.environ.get("TEST_ABC")
+
+
+assert ray.get(f.remote()) == "1"
+assert ray.get(g.remote()) is None

python/ray/tests/runtime_env_container/test_serve_basic.py

@@ -0,0 +1,41 @@
+import argparse
+
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.serve.handle import DeploymentHandle
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+args = parser.parse_args()
+
+if args.use_image_uri_api:
+    runtime_env = {"image_uri": args.image}
+else:
+    runtime_env = {"container": {"image": args.image}}
+
+
+@serve.deployment(ray_actor_options={"runtime_env": runtime_env})
+class Model:
+    def __call__(self):
+        with open("file.txt") as f:
+            return f.read().strip()
+
+
+def check_application(app_handle: DeploymentHandle, expected: str):
+    ref = app_handle.remote()
+    assert ref.result() == expected
+    return True
+
+
+h = serve.run(Model.bind())
+wait_for_condition(
+    check_application,
+    app_handle=h,
+    expected="helloworldalice",
+    timeout=300,
+)

python/ray/tests/runtime_env_container/test_serve_telemetry.py

@@ -0,0 +1,138 @@
+import argparse
+import os
+import subprocess
+
+import ray
+from ray import serve
+from ray._common.test_utils import wait_for_condition
+from ray.serve._private.test_utils import (
+    TelemetryStorage,
+    check_ray_started,
+    check_ray_stopped,
+)
+from ray.serve._private.usage import ServeUsageTag
+from ray.serve.context import _get_global_client
+from ray.serve.schema import ServeDeploySchema
+
+parser = argparse.ArgumentParser(
+    description="Example Python script taking command line arguments."
+)
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+args = parser.parse_args()
+
+os.environ["RAY_USAGE_STATS_ENABLED"] = "1"
+os.environ["RAY_USAGE_STATS_REPORT_URL"] = "http://127.0.0.1:8000/telemetry"
+os.environ["RAY_USAGE_STATS_REPORT_INTERVAL_S"] = "1"
+
+
+if args.use_image_uri_api:
+    runtime_env = {"image_uri": args.image}
+else:
+    runtime_env = {"container": {"image": args.image}}
+
+
+def check_app(app_name: str, expected: str):
+    app_handle = serve.get_app_handle(app_name)
+    ref = app_handle.remote()
+    assert ref.result() == expected
+    return True
+
+
+def check_telemetry_app():
+    report = ray.get(storage_handle.get_report.remote())
+    print(report["extra_usage_tags"])
+    assert (
+        ServeUsageTag.APP_CONTAINER_RUNTIME_ENV_USED.get_value_from_report(report)
+        == "1"
+    )
+    return True
+
+
+def check_telemetry_deployment():
+    report = ray.get(storage_handle.get_report.remote())
+    print(report["extra_usage_tags"])
+    assert (
+        ServeUsageTag.DEPLOYMENT_CONTAINER_RUNTIME_ENV_USED.get_value_from_report(
+            report
+        )
+        == "1"
+    )
+    return True
+
+
+subprocess.check_output(["ray", "start", "--head"])
+wait_for_condition(check_ray_started, timeout=5)
+serve.start()
+
+# Start TelemetryStorage and perform initial checks
+storage_handle = TelemetryStorage.remote()
+client = _get_global_client()
+config = {
+    "applications": [
+        {
+            "name": "telemetry",
+            "route_prefix": "/telemetry",
+            "import_path": "ray.serve._private.test_utils.receiver_app",
+        },
+    ],
+}
+client.deploy_apps(ServeDeploySchema.parse_obj(config))
+wait_for_condition(
+    lambda: ray.get(storage_handle.get_reports_received.remote()) > 0, timeout=5
+)
+report = ray.get(storage_handle.get_report.remote())
+assert (
+    ServeUsageTag.APP_CONTAINER_RUNTIME_ENV_USED.get_value_from_report(report) is None
+)
+assert (
+    ServeUsageTag.DEPLOYMENT_CONTAINER_RUNTIME_ENV_USED.get_value_from_report(report)
+    is None
+)
+
+# Deploy with container runtime env set at application level
+config["applications"].append(
+    {
+        "name": "app1",
+        "import_path": "serve_application:app",
+        "route_prefix": "/app1",
+        "runtime_env": runtime_env,
+    },
+)
+client.deploy_apps(ServeDeploySchema.parse_obj(config))
+wait_for_condition(check_app, app_name="app1", expected="helloworldalice", timeout=300)
+wait_for_condition(check_telemetry_app)
+
+deployment_runtime_env = dict(runtime_env)
+deployment_runtime_env["working_dir"] = None
+# Deploy with container runtime env set at deployment level
+config["applications"].append(
+    {
+        "name": "app2",
+        "import_path": "read_file:app",
+        "route_prefix": "/app2",
+        "runtime_env": {
+            "working_dir": "https://github.com/ray-project/test_dag/archive/4d2c9a59d9eabfd4c8a9e04a7aae44fc8f5b416f.zip"  # noqa
+        },
+        "deployments": [
+            {
+                "name": "Model",
+                "ray_actor_options": {
+                    "runtime_env": deployment_runtime_env,
+                },
+            }
+        ],
+    }
+)
+client.deploy_apps(ServeDeploySchema.parse_obj(config))
+wait_for_condition(check_app, app_name="app2", expected="helloworldalice", timeout=300)
+wait_for_condition(check_telemetry_deployment)
+print("Telemetry checks passed!")
+
+# Stop ray
+subprocess.check_output(["ray", "stop", "--force"])
+wait_for_condition(check_ray_stopped, timeout=15)

python/ray/tests/runtime_env_container/test_shared_memory.py

@@ -0,0 +1,35 @@
+import argparse
+import sys
+
+import numpy as np
+
+import ray
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+args = parser.parse_args()
+
+if args.use_image_uri_api:
+    runtime_env = {"image_uri": args.image}
+else:
+    runtime_env = {"container": {"image": args.image}}
+
+
+@ray.remote(runtime_env=runtime_env)
+def f():
+    array = np.random.rand(5000, 5000)
+    return ray.put(array)
+
+
+ray.init()
+ref = ray.get(f.remote())
+val = ray.get(ref)
+size = sys.getsizeof(val)
+assert size < sys.getsizeof(np.random.rand(5000, 5000))
+print(f"Size of result fetched from ray.put: {size}")
+assert val.shape == (5000, 5000)

python/ray/tests/runtime_env_container/test_with_env_vars.py

@@ -0,0 +1,40 @@
+import argparse
+import os
+
+import ray
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+args = parser.parse_args()
+
+if args.use_image_uri_api:
+    runtime_env = {
+        "image_uri": args.image,
+        "env_vars": {"TEST_DEF": "1"},
+    }
+else:
+    runtime_env = {
+        "container": {"image": args.image},
+        "env_vars": {"TEST_DEF": "hi world"},
+    }
+
+
+@ray.remote(runtime_env=runtime_env)
+def f(env_var_name: str):
+    return os.environ.get(env_var_name)
+
+
+os.environ["TEST_SCRIPT_ENV_VAR"] = "hi from driver"
+
+# Set in runtime environment `env_vars`, should be picked up
+assert ray.get(f.remote("TEST_DEF")) == "hi world"
+# Environment variables that start with prefix "RAY_" should be
+# inherited from host environment
+assert ray.get(f.remote("RAY_TEST_ABC")) == "1"
+# Environment variable from driver should not be inherited
+assert not ray.get(f.remote("TEST_SCRIPT_ENV_VAR"))

python/ray/tests/runtime_env_container/test_worker_exit_intended_system_exit_and_user_error.py

@@ -0,0 +1,148 @@
+import argparse
+import asyncio
+import os
+
+import ray
+from ray._common.test_utils import wait_for_condition
+from ray._private.state_api_test_utils import verify_failed_task
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from ray.util.state import list_workers
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--image", type=str, help="The docker image to use for Ray worker")
+parser.add_argument(
+    "--use-image-uri-api",
+    action="store_true",
+    help="Whether to use the new `image_uri` API instead of the old `container` API.",
+)
+args = parser.parse_args()
+
+
+if args.use_image_uri_api:
+    runtime_env = {"image_uri": args.image}
+else:
+    runtime_env = {"container": {"image": args.image}}
+
+ray.init(num_cpus=1)
+
+
+def get_worker_by_pid(pid, detail=True):
+    for w in list_workers(detail=detail):
+        if w["pid"] == pid:
+            return w
+    assert False
+
+
+@ray.remote(runtime_env=runtime_env)
+def f():
+    return ray.get(g.remote())
+
+
+@ray.remote(runtime_env=runtime_env)
+def g():
+    return os.getpid()
+
+
+# Start a task that has a blocking call ray.get with g.remote.
+# g.remote will borrow the CPU and start a new worker.
+# The worker started for g.remote will exit by IDLE timeout.
+pid = ray.get(f.remote())
+
+
+def verify_exit_by_idle_timeout():
+    worker = get_worker_by_pid(pid)
+    type = worker["exit_type"]
+    detail = worker["exit_detail"]
+    return type == "INTENDED_SYSTEM_EXIT" and "it was idle" in detail
+
+
+wait_for_condition(verify_exit_by_idle_timeout)
+
+ray.shutdown()
+
+
+@ray.remote(
+    num_cpus=1,
+    runtime_env=runtime_env,
+)
+class A:
+    def __init__(self):
+        self.sleeping = False
+
+    async def getpid(self):
+        while not self.sleeping:
+            await asyncio.sleep(0.1)
+        return os.getpid()
+
+    async def sleep(self):
+        self.sleeping = True
+        await asyncio.sleep(9999)
+
+
+pg = ray.util.placement_group(bundles=[{"CPU": 1}])
+a = A.options(
+    scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
+).remote()
+a.sleep.options(name="sleep").remote()
+pid = ray.get(a.getpid.remote())
+ray.util.remove_placement_group(pg)
+
+
+def verify_exit_by_pg_removed():
+    worker = get_worker_by_pid(pid)
+    type = worker["exit_type"]
+    detail = worker["exit_detail"]
+    assert verify_failed_task(
+        name="sleep",
+        error_type="ACTOR_DIED",
+        error_message=["INTENDED_SYSTEM_EXIT", "placement group was removed"],
+    )
+    return type == "INTENDED_SYSTEM_EXIT" and "placement group was removed" in detail
+
+
+wait_for_condition(verify_exit_by_pg_removed)
+
+
+@ray.remote(runtime_env=runtime_env)
+class PidDB:
+    def __init__(self):
+        self.pid = None
+
+    def record_pid(self, pid):
+        self.pid = pid
+
+    def get_pid(self):
+        return self.pid
+
+
+p = PidDB.remote()
+
+
+@ray.remote(runtime_env=runtime_env)
+class FaultyActor:
+    def __init__(self):
+        p.record_pid.remote(os.getpid())
+        raise Exception("exception in the initialization method")
+
+    def ready(self):
+        pass
+
+
+a = FaultyActor.remote()
+wait_for_condition(lambda: ray.get(p.get_pid.remote()) is not None)
+pid = ray.get(p.get_pid.remote())
+
+
+def verify_exit_by_actor_init_failure():
+    worker = get_worker_by_pid(pid)
+    type = worker["exit_type"]
+    detail = worker["exit_detail"]
+    assert type == "USER_ERROR" and "exception in the initialization method" in detail
+    return verify_failed_task(
+        name="FaultyActor.__init__",
+        error_type="TASK_EXECUTION_EXCEPTION",
+        error_message="exception in the initialization method",
+    )
+
+
+wait_for_condition(verify_exit_by_actor_init_failure)