modal 1.0.2.dev7__py3-none-any.whl → 1.0.3__py3-none-any.whl

modal/mount.py

@@ -21,7 +21,7 @@ from modal_version import __version__
 
 from ._object import _get_environment_name, _Object
 from ._resolver import Resolver
-from ._utils.async_utils import aclosing, async_map, synchronize_api
+from ._utils.async_utils import TaskContext, aclosing, async_map, synchronize_api
 from ._utils.blob_utils import FileUploadSpec, blob_upload_file, get_file_upload_spec_from_path
 from ._utils.deprecation import deprecation_warning
 from ._utils.grpc_utils import retry_transient_errors
@@ -115,7 +115,8 @@ class _MountFile(_MountEntry):
     def get_files_to_upload(self):
         local_file = self.local_file.resolve()
         if not local_file.exists():
-            raise FileNotFoundError(local_file)
+            msg = f"local file {local_file} does not exist"
+            raise FileNotFoundError(msg)
 
         rel_filename = self.remote_path
         yield local_file, rel_filename
@@ -144,10 +145,12 @@ class _MountDir(_MountEntry):
         local_dir = self.local_dir.expanduser().absolute()
 
         if not local_dir.exists():
-            raise FileNotFoundError(local_dir)
+            msg = f"local dir {local_dir} does not exist"
+            raise FileNotFoundError(msg)
 
         if not local_dir.is_dir():
-            raise NotADirectoryError(local_dir)
+            msg = f"local dir {local_dir} is not a directory"
+            raise NotADirectoryError(msg)
 
         if self.recursive:
             gen = (os.path.join(root, name) for root, dirs, files in os.walk(local_dir) for name in files)
@@ -810,3 +813,124 @@ def _is_modal_path(remote_path: PurePosixPath):
         if is_modal_path:
             return True
     return False
+
+
+REMOTE_PACKAGES_PATH = "/__modal/deps"
+REMOTE_SITECUSTOMIZE_PATH = "/pkg/sitecustomize.py"
+
+SITECUSTOMIZE_CONTENT = f"""
+# This file is automatically generated by Modal.
+# It ensures that Modal's python dependencies are available in the Python PATH,
+# while prioritizing user-installed packages.
+import sys; sys.path.append('{REMOTE_PACKAGES_PATH}')
+""".strip()
+
+
+async def _create_single_mount(
+    client: _Client,
+    builder_version: str,
+    python_version: str,
+    platform: str,
+    arch: str,
+    uv_python_platform: str = None,
+    check_if_exists: bool = True,
+):
+    import subprocess
+    import tempfile
+
+    profile_environment = config.get("environment")
+    abi_tag = "cp" + python_version.replace(".", "")
+    mount_name = f"{builder_version}-{abi_tag}-{platform}-{arch}"
+    uv_python_platform = uv_python_platform or f"{arch}-{platform}"
+
+    if check_if_exists:
+        try:
+            await Mount.from_name(mount_name, namespace=api_pb2.DEPLOYMENT_NAMESPACE_GLOBAL).hydrate.aio(client)
+            print(f"✅ Found existing mount {mount_name} in global namespace.")
+            return
+        except modal.exception.NotFoundError:
+            pass
+
+    with tempfile.TemporaryDirectory() as tmpd:
+        print(f"📦 Building {mount_name}.")
+        requirements = os.path.join(os.path.dirname(__file__), f"requirements/{builder_version}.txt")
+        subprocess.run(
+            [
+                "uv",
+                "pip",
+                "install",
+                "--strict",
+                "--no-deps",
+                "--no-cache",
+                "-r",
+                requirements,
+                "--compile-bytecode",
+                "--target",
+                tmpd,
+                "--python-platform",
+                uv_python_platform,
+                "--python-version",
+                python_version,
+            ],
+            check=True,
+            capture_output=True,
+        )
+
+        print(f"🌐 Downloaded and unpacked packages to {tmpd}.")
+
+        python_mount = Mount._from_local_dir(tmpd, remote_path=REMOTE_PACKAGES_PATH)
+
+        with tempfile.NamedTemporaryFile() as sitecustomize:
+            sitecustomize.write(
+                SITECUSTOMIZE_CONTENT.encode("utf-8"),
+            )
+            sitecustomize.flush()
+
+            python_mount = python_mount.add_local_file(
+                sitecustomize.name,
+                remote_path=REMOTE_SITECUSTOMIZE_PATH,
+            )
+
+            await python_mount._deploy.aio(
+                mount_name,
+                api_pb2.DEPLOYMENT_NAMESPACE_GLOBAL,
+                environment_name=profile_environment,
+                client=client,
+            )
+            print(f"✅ Deployed mount {mount_name} to global namespace.")
+
+
+async def _create_client_dependency_mounts(
+    client=None,
+    check_if_exists=True,
+    python_versions: list[str] = list(PYTHON_STANDALONE_VERSIONS),
+):
+    coros = []
+    for python_version in python_versions:
+        # glibc >= 2.17
+        coros.append(
+            _create_single_mount(
+                client,
+                "PREVIEW",
+                python_version,
+                "manylinux_2_17",
+                "x86_64",
+                check_if_exists=check_if_exists,
+            )
+        )
+        # musl >= 1.2
+        coros.append(
+            _create_single_mount(
+                client,
+                "PREVIEW",
+                python_version,
+                "musllinux_1_2",
+                "x86_64",
+                uv_python_platform="x86_64-unknown-linux-musl",
+                check_if_exists=check_if_exists,
+            )
+        )
+    await TaskContext.gather(*coros)
+
+
+create_client_dependency_mounts = synchronize_api(_create_client_dependency_mounts)

modal/mount.pyi

@@ -308,6 +308,28 @@ def _create_client_mount(): ...
 def create_client_mount(): ...
 def _get_client_mount(): ...
 def _is_modal_path(remote_path: pathlib.PurePosixPath): ...
+async def _create_single_mount(
+    client: modal.client._Client,
+    builder_version: str,
+    python_version: str,
+    platform: str,
+    arch: str,
+    uv_python_platform: str = None,
+    check_if_exists: bool = True,
+): ...
+async def _create_client_dependency_mounts(
+    client=None, check_if_exists=True, python_versions: list[str] = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+): ...
+
+class __create_client_dependency_mounts_spec(typing_extensions.Protocol):
+    def __call__(
+        self, /, client=None, check_if_exists=True, python_versions: list[str] = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+    ): ...
+    async def aio(
+        self, /, client=None, check_if_exists=True, python_versions: list[str] = ["3.9", "3.10", "3.11", "3.12", "3.13"]
+    ): ...
+
+create_client_dependency_mounts: __create_client_dependency_mounts_spec
 
 ROOT_DIR: pathlib.PurePosixPath
 

modal/parallel_map.py

@@ -3,6 +3,7 @@ import asyncio
 import enum
 import time
 import typing
+from asyncio import FIRST_COMPLETED
 from dataclasses import dataclass
 from typing import Any, Callable, Optional
 
@@ -110,6 +111,7 @@ async def _map_invocation(
     max_inputs_outstanding = response.max_inputs_outstanding or MAX_INPUTS_OUTSTANDING_DEFAULT
 
     have_all_inputs = False
+    map_done_event = asyncio.Event()
     inputs_created = 0
     inputs_sent = 0
     inputs_retried = 0
@@ -122,10 +124,6 @@ async def _map_invocation(
     stale_retry_duplicates = 0
     no_context_duplicates = 0
 
-    def count_update():
-        if count_update_callback is not None:
-            count_update_callback(outputs_completed, inputs_created)
-
     retry_queue = TimestampPriorityQueue()
     completed_outputs: set[str] = set()  # Set of input_ids whose outputs are complete (expecting no more values)
     input_queue: asyncio.Queue[api_pb2.FunctionPutInputsItem | None] = asyncio.Queue()
@@ -134,9 +132,8 @@ async def _map_invocation(
     )
 
     async def create_input(argskwargs):
-        nonlocal inputs_created
         idx = inputs_created
-        inputs_created += 1
+        update_state(set_inputs_created=inputs_created + 1)
         (args, kwargs) = argskwargs
         return await _create_input(args, kwargs, client.stub, idx=idx, method_name=function._use_method_name)
 
@@ -147,9 +144,27 @@ async def _map_invocation(
                 break
             yield raw_input  # args, kwargs
 
-    async def drain_input_generator():
-        nonlocal have_all_inputs
+    def update_state(set_have_all_inputs=None, set_inputs_created=None, set_outputs_completed=None):
+        # This should be the only method that needs nonlocal of the following vars
+        nonlocal have_all_inputs, inputs_created, outputs_completed
+        assert set_have_all_inputs is not False  # not allowed
+        assert set_inputs_created is None or set_inputs_created > inputs_created
+        assert set_outputs_completed is None or set_outputs_completed > outputs_completed
+        if set_have_all_inputs is not None:
+            have_all_inputs = set_have_all_inputs
+        if set_inputs_created is not None:
+            inputs_created = set_inputs_created
+        if set_outputs_completed is not None:
+            outputs_completed = set_outputs_completed
+
+        if count_update_callback is not None:
+            count_update_callback(outputs_completed, inputs_created)
+
+        if have_all_inputs and outputs_completed >= inputs_created:
+            # map is done
+            map_done_event.set()
 
+    async def drain_input_generator():
         # Parallelize uploading blobs
         async with aclosing(
             async_map_ordered(input_iter(), create_input, concurrency=BLOB_MAX_PARALLELISM)
@@ -159,12 +174,12 @@ async def _map_invocation(
 
         # close queue iterator
         await input_queue.put(None)
-        have_all_inputs = True
+        update_state(set_have_all_inputs=True)
         yield
 
     async def pump_inputs():
         assert client.stub
-        nonlocal inputs_created, inputs_sent
+        nonlocal inputs_sent
         async for items in queue_batch_iterator(input_queue, max_batch_size=MAP_INVOCATION_CHUNK_SIZE):
             # Add items to the manager. Their state will be SENDING.
             await map_items_manager.add_items(items)
@@ -178,7 +193,6 @@ async def _map_invocation(
             )
 
             resp = await send_inputs(client.stub.FunctionPutInputs, request)
-            count_update()
             inputs_sent += len(items)
             # Change item state to WAITING_FOR_OUTPUT, and set the input_id and input_jwt which are in the response.
             map_items_manager.handle_put_inputs_response(resp.inputs)
@@ -231,11 +245,8 @@ async def _map_invocation(
     async def get_all_outputs():
         assert client.stub
         nonlocal \
-            inputs_created, \
             successful_completions, \
             failed_completions, \
-            outputs_completed, \
-            have_all_inputs, \
             outputs_received, \
             already_complete_duplicates, \
             no_context_duplicates, \
@@ -244,7 +255,7 @@ async def _map_invocation(
 
         last_entry_id = "0-0"
 
-        while not have_all_inputs or outputs_completed < inputs_created:
+        while not map_done_event.is_set():
             logger.debug(f"Requesting outputs. Have {outputs_completed} outputs, {inputs_created} inputs.")
             # Get input_jwts of all items in the WAITING_FOR_OUTPUT state.
             # The server uses these to track for lost inputs.
@@ -258,12 +269,26 @@ async def _map_invocation(
                 requested_at=time.time(),
                 input_jwts=input_jwts,
             )
-            response = await retry_transient_errors(
-                client.stub.FunctionGetOutputs,
-                request,
-                max_retries=20,
-                attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
+            get_response_task = asyncio.create_task(
+                retry_transient_errors(
+                    client.stub.FunctionGetOutputs,
+                    request,
+                    max_retries=20,
+                    attempt_timeout=OUTPUTS_TIMEOUT + ATTEMPT_TIMEOUT_GRACE_PERIOD,
+                )
             )
+            map_done_task = asyncio.create_task(map_done_event.wait())
+            done, pending = await asyncio.wait([get_response_task, map_done_task], return_when=FIRST_COMPLETED)
+            if get_response_task in done:
+                map_done_task.cancel()
+                response = get_response_task.result()
+            else:
+                assert map_done_event.is_set()
+                # map is done, cancel the pending call
+                get_response_task.cancel()
+                # not strictly necessary - don't leave dangling task
+                await asyncio.gather(get_response_task, return_exceptions=True)
+                return
 
             last_entry_id = response.last_entry_id
             now_seconds = int(time.time())
@@ -288,7 +313,7 @@ async def _map_invocation(
 
                 if output_type == _OutputType.SUCCESSFUL_COMPLETION or output_type == _OutputType.FAILED_COMPLETION:
                     completed_outputs.add(item.input_id)
-                    outputs_completed += 1
+                    update_state(set_outputs_completed=outputs_completed + 1)
                     yield item
 
     async def get_all_outputs_and_clean_up():
@@ -328,7 +353,6 @@ async def _map_invocation(
             async_map_ordered(get_all_outputs_and_clean_up(), fetch_output, concurrency=BLOB_MAX_PARALLELISM)
         ) as streamer:
             async for idx, output in streamer:
-                count_update()
                 if not order_outputs:
                     yield _OutputValue(output)
                 else:
@@ -401,7 +425,7 @@ async def _map_helper(
     """
 
     raw_input_queue: Any = SynchronizedQueue()  # type: ignore
-    raw_input_queue.init()
+    await raw_input_queue.init.aio()
 
     async def feed_queue():
         async with aclosing(async_input_gen) as streamer:

modal/runner.py

@@ -9,7 +9,6 @@ import dataclasses
 import os
 import time
 import typing
-import warnings
 from collections.abc import AsyncGenerator
 from multiprocessing.synchronize import Event
 from typing import TYPE_CHECKING, Any, Optional, TypeVar
@@ -296,12 +295,8 @@ async def _run_app(
 
     output_mgr = _get_output_manager()
     if interactive and output_mgr is None:
-        warnings.warn(
-            "Interactive mode is disabled because no output manager is active. "
-            "Use 'with modal.enable_output():' to enable interactive mode and see logs.",
-            stacklevel=2,
-        )
-        interactive = False
+        msg = "Interactive mode requires output to be enabled. (Use the the `modal.enable_output()` context manager.)"
+        raise InvalidError(msg)
 
     running_app: RunningApp = await _init_local_app_new(
         client,

modal/sandbox.py

@@ -100,6 +100,7 @@ class _Sandbox(_Object, type_prefix="sb"):
         volumes: dict[Union[str, os.PathLike], Union[_Volume, _CloudBucketMount]] = {},
         pty_info: Optional[api_pb2.PTYInfo] = None,
         encrypted_ports: Sequence[int] = [],
+        h2_ports: Sequence[int] = [],
         unencrypted_ports: Sequence[int] = [],
         proxy: Optional[_Proxy] = None,
         _experimental_scheduler_placement: Optional[SchedulerPlacement] = None,
@@ -155,6 +156,12 @@ class _Sandbox(_Object, type_prefix="sb"):
 
             open_ports = [api_pb2.PortSpec(port=port, unencrypted=False) for port in encrypted_ports]
             open_ports.extend([api_pb2.PortSpec(port=port, unencrypted=True) for port in unencrypted_ports])
+            open_ports.extend(
+                [
+                    api_pb2.PortSpec(port=port, unencrypted=False, tunnel_type=api_pb2.TUNNEL_TYPE_H2)
+                    for port in h2_ports
+                ]
+            )
 
             if block_network:
                 # If the network is blocked, cidr_allowlist is invalid as we don't allow any network access.
@@ -240,6 +247,8 @@ class _Sandbox(_Object, type_prefix="sb"):
         pty_info: Optional[api_pb2.PTYInfo] = None,
         # List of ports to tunnel into the sandbox. Encrypted ports are tunneled with TLS.
         encrypted_ports: Sequence[int] = [],
+        # List of encrypted ports to tunnel into the sandbox, using HTTP/2.
+        h2_ports: Sequence[int] = [],
         # List of ports to tunnel into the sandbox without encryption.
         unencrypted_ports: Sequence[int] = [],
         # Reference to a Modal Proxy to use in front of this Sandbox.
@@ -283,6 +292,7 @@ class _Sandbox(_Object, type_prefix="sb"):
             volumes=volumes,
             pty_info=pty_info,
             encrypted_ports=encrypted_ports,
+            h2_ports=h2_ports,
             unencrypted_ports=unencrypted_ports,
             proxy=proxy,
             _experimental_enable_snapshot=_experimental_enable_snapshot,
@@ -320,6 +330,8 @@ class _Sandbox(_Object, type_prefix="sb"):
         pty_info: Optional[api_pb2.PTYInfo] = None,
         # List of ports to tunnel into the sandbox. Encrypted ports are tunneled with TLS.
         encrypted_ports: Sequence[int] = [],
+        # List of encrypted ports to tunnel into the sandbox, using HTTP/2.
+        h2_ports: Sequence[int] = [],
         # List of ports to tunnel into the sandbox without encryption.
         unencrypted_ports: Sequence[int] = [],
         # Reference to a Modal Proxy to use in front of this Sandbox.
@@ -359,6 +371,7 @@ class _Sandbox(_Object, type_prefix="sb"):
             volumes=volumes,
             pty_info=pty_info,
             encrypted_ports=encrypted_ports,
+            h2_ports=h2_ports,
             unencrypted_ports=unencrypted_ports,
             proxy=proxy,
             _experimental_scheduler_placement=_experimental_scheduler_placement,
@@ -517,7 +530,7 @@ class _Sandbox(_Object, type_prefix="sb"):
 
         return self._tunnels
 
-    async def terminate(self):
+    async def terminate(self) -> None:
         """Terminate Sandbox execution.
 
         This is a no-op if the Sandbox has already finished running."""
@@ -525,7 +538,6 @@ class _Sandbox(_Object, type_prefix="sb"):
         await retry_transient_errors(
             self._client.stub.SandboxTerminate, api_pb2.SandboxTerminateRequest(sandbox_id=self.object_id)
         )
-        await self.wait(raise_on_termination=False)
 
     async def poll(self) -> Optional[int]:
         """Check if the Sandbox has finished running.
@@ -541,7 +553,7 @@ class _Sandbox(_Object, type_prefix="sb"):
 
         return self.returncode
 
-    async def _get_task_id(self):
+    async def _get_task_id(self) -> str:
         while not self._task_id:
             resp = await self._client.stub.SandboxGetTaskId(api_pb2.SandboxGetTaskIdRequest(sandbox_id=self.object_id))
             self._task_id = resp.task_id

modal/sandbox.pyi

@@ -58,6 +58,7 @@ class _Sandbox(modal._object._Object):
         ] = {},
         pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
         encrypted_ports: collections.abc.Sequence[int] = [],
+        h2_ports: collections.abc.Sequence[int] = [],
         unencrypted_ports: collections.abc.Sequence[int] = [],
         proxy: typing.Optional[modal.proxy._Proxy] = None,
         _experimental_scheduler_placement: typing.Optional[modal.scheduler_placement.SchedulerPlacement] = None,
@@ -86,6 +87,7 @@ class _Sandbox(modal._object._Object):
         ] = {},
         pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
         encrypted_ports: collections.abc.Sequence[int] = [],
+        h2_ports: collections.abc.Sequence[int] = [],
         unencrypted_ports: collections.abc.Sequence[int] = [],
         proxy: typing.Optional[modal.proxy._Proxy] = None,
         _experimental_enable_snapshot: bool = False,
@@ -116,6 +118,7 @@ class _Sandbox(modal._object._Object):
         ] = {},
         pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
         encrypted_ports: collections.abc.Sequence[int] = [],
+        h2_ports: collections.abc.Sequence[int] = [],
         unencrypted_ports: collections.abc.Sequence[int] = [],
         proxy: typing.Optional[modal.proxy._Proxy] = None,
         _experimental_enable_snapshot: bool = False,
@@ -129,9 +132,9 @@ class _Sandbox(modal._object._Object):
     async def snapshot_filesystem(self, timeout: int = 55) -> modal.image._Image: ...
     async def wait(self, raise_on_termination: bool = True): ...
     async def tunnels(self, timeout: int = 50) -> dict[int, modal._tunnel.Tunnel]: ...
-    async def terminate(self): ...
+    async def terminate(self) -> None: ...
     async def poll(self) -> typing.Optional[int]: ...
-    async def _get_task_id(self): ...
+    async def _get_task_id(self) -> str: ...
     @typing.overload
     async def exec(
         self,
@@ -228,6 +231,7 @@ class Sandbox(modal.object.Object):
         ] = {},
         pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
         encrypted_ports: collections.abc.Sequence[int] = [],
+        h2_ports: collections.abc.Sequence[int] = [],
         unencrypted_ports: collections.abc.Sequence[int] = [],
         proxy: typing.Optional[modal.proxy.Proxy] = None,
         _experimental_scheduler_placement: typing.Optional[modal.scheduler_placement.SchedulerPlacement] = None,
@@ -261,6 +265,7 @@ class Sandbox(modal.object.Object):
             ] = {},
             pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
             encrypted_ports: collections.abc.Sequence[int] = [],
+            h2_ports: collections.abc.Sequence[int] = [],
             unencrypted_ports: collections.abc.Sequence[int] = [],
             proxy: typing.Optional[modal.proxy.Proxy] = None,
             _experimental_enable_snapshot: bool = False,
@@ -293,6 +298,7 @@ class Sandbox(modal.object.Object):
             ] = {},
             pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
             encrypted_ports: collections.abc.Sequence[int] = [],
+            h2_ports: collections.abc.Sequence[int] = [],
             unencrypted_ports: collections.abc.Sequence[int] = [],
             proxy: typing.Optional[modal.proxy.Proxy] = None,
             _experimental_enable_snapshot: bool = False,
@@ -330,6 +336,7 @@ class Sandbox(modal.object.Object):
             ] = {},
             pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
             encrypted_ports: collections.abc.Sequence[int] = [],
+            h2_ports: collections.abc.Sequence[int] = [],
             unencrypted_ports: collections.abc.Sequence[int] = [],
             proxy: typing.Optional[modal.proxy.Proxy] = None,
             _experimental_enable_snapshot: bool = False,
@@ -363,6 +370,7 @@ class Sandbox(modal.object.Object):
             ] = {},
             pty_info: typing.Optional[modal_proto.api_pb2.PTYInfo] = None,
             encrypted_ports: collections.abc.Sequence[int] = [],
+            h2_ports: collections.abc.Sequence[int] = [],
             unencrypted_ports: collections.abc.Sequence[int] = [],
             proxy: typing.Optional[modal.proxy.Proxy] = None,
             _experimental_enable_snapshot: bool = False,
@@ -405,8 +413,8 @@ class Sandbox(modal.object.Object):
     tunnels: __tunnels_spec[typing_extensions.Self]
 
     class __terminate_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /): ...
-        async def aio(self, /): ...
+        def __call__(self, /) -> None: ...
+        async def aio(self, /) -> None: ...
 
     terminate: __terminate_spec[typing_extensions.Self]
 
@@ -417,8 +425,8 @@ class Sandbox(modal.object.Object):
     poll: __poll_spec[typing_extensions.Self]
 
     class ___get_task_id_spec(typing_extensions.Protocol[SUPERSELF]):
-        def __call__(self, /): ...
-        async def aio(self, /): ...
+        def __call__(self, /) -> str: ...
+        async def aio(self, /) -> str: ...
 
     _get_task_id: ___get_task_id_spec[typing_extensions.Self]
 

modal/schedule.py

@@ -30,15 +30,28 @@ class Cron(Schedule):
     We can specify different schedules with cron strings, for example:
 
     ```python
-    modal.Cron("5 4 * * *")  # run at 4:05am every night
-    modal.Cron("0 9 * * 4")  # runs every Thursday 9am
+    modal.Cron("5 4 * * *")  # run at 4:05am UTC every night
+    modal.Cron("0 9 * * 4")  # runs every Thursday at 9am UTC
     ```
 
+    We can also optionally specify a timezone, for example:
+
+    ```python
+    # Run daily at 6am New York time, regardless of whether daylight saving
+    # is in effect (i.e. at 11am UTC in the winter, and 10am UTC in the summer):
+    modal.Cron("0 6 * * *", timezone="America/New_York")
+    ```
+
+    If no timezone is specified, the default is UTC.
     """
 
-    def __init__(self, cron_string: str) -> None:
+    def __init__(
+        self,
+        cron_string: str,
+        timezone: str = "UTC",
+    ) -> None:
         """Construct a schedule that runs according to a cron expression string."""
-        cron = api_pb2.Schedule.Cron(cron_string=cron_string)
+        cron = api_pb2.Schedule.Cron(cron_string=cron_string, timezone=timezone)
         super().__init__(api_pb2.Schedule(cron=cron))