import os
import logging
try:
from dask_gateway import Gateway
except ImportError:
logging.error("Unable to import Dask Gateway. Are you running in a cloud compute environment?\n")
raise
os.environ['DASK_DISTRIBUTED__SCHEDULER__WORKER_SATURATION'] = "1.0"
gateway = Gateway()
_options = gateway.cluster_options()
_options.conda_environment='global/global-pangeo' ##<< this is the conda environment we use on nebari.
_options.profile = 'Medium Worker'
_env_to_add={}
aws_env_vars=['AWS_ACCESS_KEY_ID',
'AWS_SECRET_ACCESS_KEY',
'AWS_SESSION_TOKEN',
'AWS_DEFAULT_REGION',
'AWS_S3_ENDPOINT']
for _e in aws_env_vars:
if _e in os.environ:
_env_to_add[_e] = os.environ[_e]
_options.environment_vars = _env_to_add
cluster = gateway.new_cluster(_options) ##<< create cluster via the dask gateway
cluster.adapt(minimum=2, maximum=30) ##<< Sets scaling parameters.
client = cluster.get_client()
print("The 'cluster' object can be used to adjust cluster behavior. i.e. 'cluster.adapt(minimum=10)'")
print("The 'client' object can be used to directly interact with the cluster. i.e. 'client.submit(func)' ")
print(f"The link to view the client dashboard is:\n> {client.dashboard_link}")
