rs_common/prefect_utils.md

<< Back to index

Utility Python module for the tutorials, to be shared with the prefect or dask workers.

WARNING: AFTER EACH MODIFICATION, RESTART THE JUPYTER NOTEBOOK KERNEL !

delete_payload_file(payload_path, log=False)

Delete either an S3 object or a local file.

Parameters:

Name	Type	Description	Default
payload_path	str	Path to delete (S3 path starting with s3:// or local file path)	required
log	bool	Whether to log the deletion (for S3 paths only)	False

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

def delete_payload_file(payload_path: str, log: bool = False) -> None:
    """Delete either an S3 object or a local file.

    Args:
        payload_path: Path to delete (S3 path starting with s3:// or local file path)
        log: Whether to log the deletion (for S3 paths only)
    """
    if payload_path.startswith(("s3://", "S3://")):
        # Handle S3 path
        s3_delete(payload_path, log=log)
    else:
        # Handle local file path
        try:
            if log:
                logger.info(f"Deleting local file: {payload_path}")
            os.remove(payload_path)
            if log:
                logger.info(f"Deleted local file: {payload_path}")
        except FileNotFoundError:
            # File already deleted, ignore
            pass

format_env_user(block_name, any_owner_id)

Format the Prefect secret block for the current user

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

def format_env_user(block_name: str, any_owner_id: str):
    """Format the Prefect secret block for the current user"""
    name = block_name.format(any_owner_id).lower()
    return re.sub("[^a-zA-Z0-9]", "-", name)  # replace special characters by dash

get_ip_address()

Return IP address, see: https://stackoverflow.com/a/166520

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

def get_ip_address() -> str:
    """Return IP address, see: https://stackoverflow.com/a/166520"""
    return socket.gethostbyname(socket.gethostname())

get_s3_bucket(s3_path)

Return a prefect S3 bucket object and S3 "object name" (= S3 path without s3://bucket-name) from the given S3 path. We use the prefect higher-level functions instead of those from boto3.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

def get_s3_bucket(s3_path: str) -> tuple[S3Bucket, str]:
    """
    Return a prefect S3 bucket object and S3 "object name" (= S3 path without s3://bucket-name)
    from the given S3 path.
    We use the prefect higher-level functions instead of those from boto3.
    """

    # Remove the s3:// prefix and split by /
    split = s3_path.removeprefix("s3").removeprefix("S3").strip(":/").split("/")

    # Filter empty elements (if we had double //)
    split = list(filter(None, split))

    if not split:
        raise ValueError(f"Invalid S3 path: {s3_path!r}")

    bucket_name = split[0]
    object_name = "/".join(split[1:])

    # Try to return an existing bucket object
    try:
        return S3_BUCKETS[bucket_name], object_name

    # Else create a new one
    except KeyError:
        aws_credentials = AwsCredentials(
            aws_access_key_id=os.environ["S3_ACCESSKEY"],
            aws_secret_access_key=SecretStr(os.environ["S3_SECRETKEY"]),
            region_name=os.environ["S3_REGION"],
            aws_client_parameters=AwsClientParameters(endpoint_url=os.environ["S3_ENDPOINT"]),
        )
        s3_bucket = S3Bucket(
            bucket_name=bucket_name,
            credentials=aws_credentials,
            bucket_folder="",  # no prefixed folder
        )
        S3_BUCKETS[bucket_name] = s3_bucket
        return s3_bucket, object_name

get_share_bucket(sub_folder='') async

Get the prefect share bucket folder.

Parameters:

Name	Type	Description	Default
sub_folder	str	subfolder to use in the bucket. If empty, use the default folder and secret block.	''
Else use a specific secret block. WARNING		this only works in a single-threaded environment	required

Returns:

Type	Description
tuple[S3Bucket, str]	The prefect block for the share bucket and the block name.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def get_share_bucket(sub_folder: str = "") -> tuple[S3Bucket, str]:
    """
    Get the prefect share bucket folder.

    Args:
        sub_folder: subfolder to use in the bucket. If empty, use the default folder and secret block.
        Else use a specific secret block. WARNING: this only works in a single-threaded environment
        else these specific secret blocks may overwrite each other.

    Returns:
        The prefect block for the share bucket and the block name.
    """

    # Use the default secret block for the default folder
    if not sub_folder:
        block_name = BLOCK_NAME_SHARE_BUCKET_GLOBAL

        # Try to read and return the prefect block, if it already exists.
        try:
            return await S3Bucket.load(block_name), block_name

        # If it doesn't exist yet, create and return it
        except ValueError:
            pass

    # Use a specific block for the current user. Always overwrite it, because the subfolder may have changed.
    else:
        block_name = format_env_user(BLOCK_NAME_SHARE_BUCKET_USER, OWNER_ID)  # type: ignore

    # Read the prefect blocks that contain the S3 authentication
    bucket_name = os.environ["PREFECT_BUCKET_NAME"]
    bucket_folder = os.environ["PREFECT_BUCKET_FOLDER"]

    if sub_folder:
        bucket_folder = os.path.join(bucket_folder, sub_folder)

    # Get a s3 bucket object from its name
    generic_bucket, _ = get_s3_bucket(bucket_name)

    # Create a new object with the same credentials and a prefixed folder
    share_bucket = S3Bucket(
        bucket_name=bucket_name,
        bucket_folder=bucket_folder,
        credentials=generic_bucket.credentials,
    )

    # Save it as a prefect block and return it
    await share_bucket.save(block_name, overwrite=True)
    return share_bucket, block_name

hack_for_jupyter(func, *args, **kwargs)

From Jupyter we need this hack to deploy prefect flows

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

def hack_for_jupyter(func: Callable, *args, **kwargs) -> asyncio.Task:
    """From Jupyter we need this hack to deploy prefect flows"""
    coroutine = run_in_threadpool(func, *args, **kwargs)
    return asyncio.create_task(coroutine)

init_global_env(new_owner_id=None)

Init the global variables above from the env vars. Needs to be called from the client, then called again from the prefect workers.

Parameters:

Name	Type	Description	Default
new_owner_id	str | None	new ower/user id, if known.	None

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

def init_global_env(new_owner_id: str | None = None):
    """
    Init the global variables above from the env vars.
    Needs to be called from the client, then called again from the prefect workers.

    Args:
        new_owner_id: new ower/user id, if known.
    """
    global LOCAL_MODE, CLUSTER_MODE, OWNER_ID

    LOCAL_MODE = env_bool("RSPY_LOCAL_MODE", default=False)
    CLUSTER_MODE = not LOCAL_MODE

    # Override the owner id or read it from the env vars
    if new_owner_id:
        OWNER_ID = new_owner_id
    else:
        OWNER_ID = os.getenv("JUPYTERHUB_USER", "") if CLUSTER_MODE else os.getenv("RSPY_HOST_USER", "")

init_prefect_blocks() async

Init prefect blocks and variables from the client environment (= from jupyter)

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def init_prefect_blocks():
    """Init prefect blocks and variables from the client environment (= from jupyter)"""

    # In local mode, automatically create prefect blocks and variables.
    # In cluster mode, they have already been created by an admin by using rs-demo/notebooks/init-prefect/ notebooks.
    if LOCAL_MODE:

        #
        # Env vars for all users

        # Get all env var names that start with DASK_GATEWAY_
        regex = re.compile("^DASK_GATEWAY_.*")
        dask_gateway_vars = [v for v in os.environ if regex.match(v)]

        # Copy env vars. They are all mandatory.
        global_env_vars = {}
        for key in [
            "RSPY_LOCAL_MODE",
            "RSPY_PREFECT_URL",
            "PREFECT_BUCKET_NAME",
            "PREFECT_BUCKET_FOLDER",
            "POSTGRES_USER",
            "POSTGRES_PASSWORD",
            "POSTGRES_PORT",
            "POSTGRES_PI_DB",
            "POSTGRES_HOST",
            *dask_gateway_vars,
            "LOCAL_DASK_USERNAME",
            "LOCAL_DASK_PASSWORD",
        ]:
            global_env_vars[key] = os.environ[key]

        # These ones are used only in cluster mode. In local mode we set empty values.
        for key in [
            "JUPYTERHUB_API_TOKEN",
            "DASK_GATEWAY_ADDRESS",
        ]:
            global_env_vars[key] = ""

        # Save env vars in a secret block for all users
        await update_prefect_block(BLOCK_NAME_ENV_GLOBAL, global_env_vars)

        #
        # Processing storage configuration

        if not await Variable.get(VAR_NAME_STORAGE_CONFIG):  # if this var doesn't already exist

            # Read config file content
            async with await anyio.open_file(str(CONFIG_DIR / "storage_configuration.json"), encoding="utf-8") as f:

                # Write Prefect variable
                content = await f.read()
                await Variable.set(VAR_NAME_STORAGE_CONFIG, json.loads(content))

    #
    # Env vars for current user/owner_id

    # In cluster mode, read the API key
    if CLUSTER_MODE:
        await read_apikey()

    # Read env vars that are available from the client env in both local and cluster mode.
    user_env_vars = {}

    # Read optional env vars
    for key in ["RSPY_APIKEY"]:
        if value := os.getenv(key):
            user_env_vars[key] = value

    # These are False by default. Save them only if they are set to True.
    for key in [
        "OTEL_PYTHON_REQUESTS_TRACE_HEADERS",
        "OTEL_PYTHON_REQUESTS_TRACE_BODY",
    ]:
        if env_bool(key, default=False):
            user_env_vars[key] = "1"

    # Save env vars in a secret block for the current user
    await update_prefect_block(format_env_user(BLOCK_NAME_ENV_USER, OWNER_ID), user_env_vars)

    # Now read back the blocks so we are sure our env vars are up-to-date
    await read_prefect_blocks(OWNER_ID)

read_apikey(optional=False, save_to_env=True) async

Read the API key, either from the environment variable or from an interactive input form.

Parameters:

Name	Type	Description	Default
optional	bool	If False and if the env var is missing, ask it from an interactive input form.	False
save_to_env	bool	If True, saves the API key to the ~/.env file.	True

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def read_apikey(optional: bool = False, save_to_env: bool = True) -> None:
    """
    Read the API key, either from the environment variable or from an interactive input form.

    Args:
        optional (bool): If False and if the env var is missing, ask it from an interactive input form.
        save_to_env (bool): If True, saves the API key to the ~/.env file.
    """
    # No API key in local mode
    if LOCAL_MODE:
        return

    # If the API is saved as an env var in the ~/.env file, then it has already
    # been read automatically by rs-infra-core/.github/jupyter/resources/00-read-env.py
    apikey = os.getenv("RSPY_APIKEY")
    if apikey:
        logger.debug(f"Use API key (probably from '~/.env'): '{apikey[:4]}***'")
    elif optional:
        logger.debug("Don't use any API key")
    else:
        # Else read it from user input
        apikey = getpass.getpass("Enter your API key:")

        # Save the env var
        os.environ["RSPY_APIKEY"] = apikey

        # Append it to the ~/.env file, if requested.
        # Don't overwrite the full ~/.env file because it can contain other user info.
        if save_to_env:
            await asyncio.to_thread(_append_apikey_to_env, apikey)
            logger.debug("API key saved to '~/.env'")

read_artifact(http_session, flow_run_id, artifact_key)

Read artifact value.

Parameters:

Name	Type	Description	Default
http_session	Session	HTTP request session	required
flow_run_id	str | UUID	Prefect flow run ID	required
artifact_key	str	Artifact key (=name)	required

Returns:

Type	Description
Any	Artifact value

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

def read_artifact(http_session: requests.Session, flow_run_id: str | UUID, artifact_key: str) -> Any:
    """
    Read artifact value.

    Args:
        http_session: HTTP request session
        flow_run_id: Prefect flow run ID
        artifact_key: Artifact key (=name)

    Returns:
        Artifact value
    """
    # See: https://docs.prefect.io/v3/api-ref/rest-api/server/artifacts/read-artifacts
    flow_run_id = str(flow_run_id)
    response = http_session.post(
        f"{os.environ['PREFECT_API_URL']}/artifacts/filter",
        json={
            "artifacts": {
                "operator": "and_",
                "key": {"any_": [artifact_key]},
                "flow_run_id": {"any_": [flow_run_id]},
            },
        },
        timeout=30,
    )
    response.raise_for_status()
    artifacts = response.json()
    if not artifacts:
        raise RuntimeError(f"No artifact found for flow run ID: {flow_run_id!r} and artifact key: {artifact_key!r}")
    return artifacts[0]["data"]

read_prefect_blocks(any_owner_id=None) async

Read prefect blocks from the prefect flow and tasks into env vars and global vars.

Parameters:

Name	Type	Description	Default
owner_id		Read prefect blocks for a specific user.	required

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def read_prefect_blocks(any_owner_id: str | None = None):
    """
    Read prefect blocks from the prefect flow and tasks into env vars and global vars.

    Args:
        owner_id: Read prefect blocks for a specific user.
    """

    # Read the env vars for all users
    os.environ.update((await Secret.load(BLOCK_NAME_ENV_GLOBAL)).get())

    # Read the env vars for the given user
    if any_owner_id:
        os.environ.update((await Secret.load(format_env_user(BLOCK_NAME_ENV_USER, any_owner_id))).get())

    # Init the env of the current module from the env vars we have just read
    init_global_env(any_owner_id)

remove_bucket_credentials(obs_ids) async

Remove bucket credentials from the prefect secret block that contains the env variables for the current user/owner_id.

Parameters:

Name	Type	Description	Default
obs_ids	list[str] | str	"" keys to remove from the env vars:	required

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def remove_bucket_credentials(obs_ids: list[str] | str):
    """
    Remove bucket credentials from the prefect secret block that contains
    the env variables for the current user/owner_id.

    Args:
        obs_ids: "<OBS_ID>" keys to remove from the env vars:
        S3_<OBS_ID>_ACCESSKEY, S3_<OBS_ID>_SECRETKEY, S3_<OBS_ID>_ENDPOINT, S3_<OBS_ID>_REGION
    """
    # Convert string to list
    if isinstance(obs_ids, str):
        obs_ids = [obs_ids]

    env_vars = []
    for key in obs_ids:
        key = key.upper()
        env_vars += [f"S3_{key}_ACCESSKEY", f"S3_{key}_SECRETKEY", f"S3_{key}_ENDPOINT", f"S3_{key}_REGION"]

    # Update the prefect block
    await update_prefect_block(format_env_user(BLOCK_NAME_ENV_USER, OWNER_ID), remove_keys=env_vars)

s3_copy_from_s3(from_s3_path, to_s3_path) async

Copy all objects found from one s3 path to another s3 path.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def s3_copy_from_s3(from_s3_path: str, to_s3_path: str) -> None:
    """Copy all objects found from one s3 path to another s3 path."""
    from_bucket, from_prefix = get_s3_bucket(from_s3_path)
    _, to_prefix = get_s3_bucket(to_s3_path)
    s3_client = from_bucket._get_s3_client()  # pylint: disable=protected-access
    bucket_name = from_bucket.bucket_name
    objects = s3_client.list_objects_v2(Bucket=bucket_name, Prefix=from_prefix)
    logger.info(f"Starting the copy from {from_s3_path} to {to_s3_path}")
    if "Contents" in objects:
        for obj in objects["Contents"]:
            src_key = obj["Key"]
            dst_key = src_key.replace(from_prefix, to_prefix, 1)
            copy_source = {"Bucket": bucket_name, "Key": src_key}
            s3_client.copy_object(CopySource=copy_source, Bucket=bucket_name, Key=dst_key)
    logger.info(f"Copied files from {from_s3_path} to {to_s3_path}")

s3_delete(s3_prefix, log=False)

Remove all files from S3 bucket with the given prefix, using low-level client.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

def s3_delete(s3_prefix: str, log: bool = False):
    """Remove all files from S3 bucket with the given prefix, using low-level client."""
    s3_bucket, prefix = get_s3_bucket(s3_prefix)
    objects = s3_bucket._get_bucket_resource().objects  # pylint: disable=protected-access

    # objects.filter(Prefix=prefix) makes a recursive search and returns only files, not folders.
    # If we want to delete a file, call the filter and only keep the exact matches.
    try_file = [obj.key for obj in objects.filter(Prefix=prefix)]
    if prefix in try_file:
        objects_to_delete = [{"Key": prefix}]

    # If we want to delete a folder, make sure the prefix ends by /, call the filter and keep everything
    else:
        objects_to_delete = [{"Key": obj.key} for obj in objects.filter(Prefix=prefix.rstrip("/") + "/")]

    if not objects_to_delete:
        return

    s3_client = s3_bucket._get_s3_client()  # pylint: disable=protected-access
    endpoint_url = s3_bucket.credentials.aws_client_parameters.endpoint_url

    if log:
        keys = [f"s3://{s3_bucket.bucket_name}/{o.get('Key')}" for o in objects_to_delete]
        s3_bucket.logger.info(
            f"Delete from {endpoint_url!r}: {json.dumps(keys, indent=2)}",  # nosec hardcoded_sql_expressions
        )

    # Split the list of objects to delete
    chunk_size = 1000
    for i in range(0, len(objects_to_delete), chunk_size):
        s3_client.delete_objects(
            Bucket=s3_bucket.bucket_name,
            Delete={"Objects": objects_to_delete[i : i + chunk_size], "Quiet": True},  # noqa: E203
        )

s3_download_dir(s3_path, local_path=None) async

See: S3Bucket.get_directory

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def s3_download_dir(
    s3_path: str,
    local_path: str | None = None,
) -> None:
    """See: S3Bucket.get_directory"""
    s3_bucket, from_path = get_s3_bucket(s3_path)
    await s3_bucket.aget_directory(from_path, local_path)

s3_download_file(s3_path, to_path, **download_kwargs) async

See: S3Bucket.download_object_to_path

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def s3_download_file(
    s3_path: str,
    to_path: str | Path | None,
    **download_kwargs: dict[str, Any],
) -> Path:
    """See: S3Bucket.download_object_to_path"""
    s3_bucket, from_path = get_s3_bucket(s3_path)
    return await s3_bucket.adownload_object_to_path(from_path, to_path, **download_kwargs)

s3_read_bytes(s3_path) async

Read an object's bytes straight from the S3 bucket, without a temporary file.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def s3_read_bytes(s3_path: str) -> bytes:
    """Read an object's bytes straight from the S3 bucket, without a temporary file."""
    s3_bucket, from_path = get_s3_bucket(s3_path)
    return await s3_bucket.aread_path(from_path)

s3_upload_bytes(data, s3_path, **upload_kwargs) async

Upload in-memory bytes straight to the S3 bucket, without writing a temporary file.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

async def s3_upload_bytes(data: bytes, s3_path: str, **upload_kwargs: dict[str, Any]) -> str:
    """Upload in-memory bytes straight to the S3 bucket, without writing a temporary file."""
    s3_bucket, to_path = get_s3_bucket(s3_path)
    return await s3_bucket.aupload_from_file_object(io.BytesIO(data), to_path, **upload_kwargs)

s3_upload_dir(from_folder, s3_path, **upload_kwargs) async

See: S3Bucket.upload_from_folder

Uploads files within a folder (excluding the folder itself) to the object storage service folder.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def s3_upload_dir(
    from_folder: str | Path,
    s3_path: str,
    **upload_kwargs: dict[str, Any],
) -> str | None:
    """
    See: S3Bucket.upload_from_folder

    Uploads files *within* a folder (excluding the folder itself) to the object storage service folder.
    """
    s3_bucket, to_path = get_s3_bucket(s3_path)
    logger.debug(f"Uploading folder {from_folder} to {s3_path}")
    return await s3_bucket.aupload_from_folder(from_folder, to_path, **upload_kwargs)

s3_upload_empty_file(s3_path, **upload_kwargs) async

Upload an empty temp file to the S3 bucket.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def s3_upload_empty_file(
    s3_path: str,
    **upload_kwargs: dict[str, Any],
) -> str:
    """Upload an empty temp file to the S3 bucket."""
    # Add contents to the file or boto3 has a strange behavior after uploading an empty file
    return await s3_upload_bytes(b"empty", s3_path, **upload_kwargs)

s3_upload_file(from_path, s3_path, **upload_kwargs) async

See: S3Bucket.upload_from_path

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def s3_upload_file(
    from_path: str | Path,
    s3_path: str,
    **upload_kwargs: dict[str, Any],
) -> str:
    """See: S3Bucket.upload_from_path"""
    s3_bucket, to_path = get_s3_bucket(s3_path)
    return await s3_bucket.aupload_from_path(from_path, to_path, **upload_kwargs)

save_bucket_credentials(osam_client=None, obs_credentials=None) async

Save bucket credentials in the prefect secret block that contains the env variables for the current user/owner_id.

Parameters:

Name	Type	Description	Default
osam_client	OsamClient | None	Used to read the user's credentials from the OSAM service. They are saved in the env vars:	None
obs_credentials	dict[str, BucketCredentials] | None	Used to save custom bucket credentials. For each dict key "", the credentials are	None
saved in the env vars		S3_ACCESSKEY, S3SECRETKEY, S3ENDPOINT, S3_REGION	required

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def save_bucket_credentials(
    osam_client: OsamClient | None = None,
    obs_credentials: dict[str, BucketCredentials] | None = None,
):
    """
    Save bucket credentials in the prefect secret block that contains
    the env variables for the current user/owner_id.

    Args:
        osam_client: Used to read the user's credentials from the OSAM service. They are saved in the env vars:
        osam_client, S3_SECRETKEY, S3_ENDPOINT, S3_REGION
        obs_credentials: Used to save custom bucket credentials. For each dict key "<OBS_ID>", the credentials are
        saved in the env vars: S3_<OBS_ID>_ACCESSKEY, S3_<OBS_ID>_SECRETKEY, S3_<OBS_ID>_ENDPOINT, S3_<OBS_ID>_REGION
    """
    env_vars = {}
    if obs_credentials is None:
        obs_credentials = {}

    # Read the user's credentials from osam
    if osam_client:
        credentials = await osam_client.get_credentials()
        env_vars.update(
            {
                "S3_ACCESSKEY": credentials.access_key,
                "S3_SECRETKEY": credentials.secret_key,
                "S3_ENDPOINT": credentials.endpoint,
                "S3_REGION": credentials.region,
            },
        )

    # Add custom credentials
    for key, credentials in obs_credentials.items():
        key = key.upper()
        env_vars.update(
            {
                f"S3_{key}_ACCESSKEY": credentials.access_key,
                f"S3_{key}_SECRETKEY": credentials.secret_key,
                f"S3_{key}_ENDPOINT": credentials.endpoint,
                f"S3_{key}_REGION": credentials.region,
            },
        )

    # Update the prefect block
    await update_prefect_block(format_env_user(BLOCK_NAME_ENV_USER, OWNER_ID), env_vars)

    # Save the env vars
    os.environ.update(env_vars)

update_prefect_block(block_name, add_values=None, remove_keys=None) async

Write a prefect secret block with values, or update it if it already exists.

Parameters:

Name	Type	Description	Default
block_name	str	Prefect secret block name	required
add_values	dict | None	Key/values to add from the block	None
remove_keys	list | None	Key/values to remove from the block	None

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

@sync_compatible
async def update_prefect_block(block_name: str, add_values: dict | None = None, remove_keys: list | None = None):
    """
    Write a prefect secret block with values, or update it if it already exists.

    Args:
        block_name: Prefect secret block name
        add_values: Key/values to add from the block
        remove_keys: Key/values to remove from the block
    """
    if add_values is None:
        add_values = {}
    if remove_keys is None:
        remove_keys = []

    with lock:
        # Try to read the block, if it exists
        try:
            old_values = (await Secret.load(block_name)).get()
        except ValueError:
            old_values = {}

        # Add new values
        updated_values = old_values | add_values

        # Remove values
        for key in remove_keys:
            updated_values.pop(key, None)

        # Write the block if it didn't exist or if values have changed
        if (not old_values) or (old_values != updated_values):
            await Secret(value=updated_values).save(block_name, overwrite=True)

upload_payload_bytes(data, payload_path, **upload_kwargs) async

Upload in-memory bytes to either S3 bucket or local file.

Parameters:

Name	Type	Description	Default
data	bytes	The data to upload as bytes	required
payload_path	str	Destination path (S3 path starting with s3:// or local file path)	required
upload_kwargs	dict[str, Any]	Additional keyword arguments for S3 upload	{}

Returns:

Type	Description
	The payload_path that was written to

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

async def upload_payload_bytes(data: bytes, payload_path: str, **upload_kwargs: dict[str, Any]):
    """Upload in-memory bytes to either S3 bucket or local file.

    Args:
        data: The data to upload as bytes
        payload_path: Destination path (S3 path starting with s3:// or local file path)
        upload_kwargs: Additional keyword arguments for S3 upload

    Returns:
        The payload_path that was written to
    """
    if payload_path.startswith(("s3://", "S3://")):
        # Handle S3 path
        await s3_upload_bytes(data, payload_path, **upload_kwargs)
    else:
        # Handle local file path
        # Ensure the directory exists
        os.makedirs(os.path.dirname(payload_path), exist_ok=True)
        # Write the data to the local file using async file I/O
        async with await anyio.open_file(payload_path, "wb") as f:
            await f.write(data)

wait_for_deployment(name, wait=1, max_retry=30) async

Wait for prefect deployment to be finished.

Source code in docs/rs-client-libraries/rs_common/prefect_utils.py

async def wait_for_deployment(name: str, wait: int | float = 1, max_retry: int = 30):
    """Wait for prefect deployment to be finished."""
    # Taken from prefect/cli/deployment.py::inspect
    retry = 0
    async with get_client() as client:
        while True:
            try:
                await client.read_deployment_by_name(name)
                logger.info(f"Finished deploying prefect flow: {name!r}")
                return
            except ObjectNotFound:
                retry += 1
                if retry >= max_retry:
                    raise
                logger.info(f"Wait for deployment of prefect flow: {name!r} ...")
                await asyncio.sleep(wait)