def test_azure_serialize_with_flows(): azure = storage.Azure( container="container", connection_string="conn", blob_name="name", secrets=["foo"], )
def test_azure_empty_serialize(): azure = storage.Azure(container="container") serialized = AzureSchema().dump(azure) assert serialized assert serialized["__version__"] == prefect.__version__ assert serialized["container"] == "container" assert serialized["blob_name"] is None assert serialized["secrets"] == []
def test_azure_full_serialize(): azure = storage.Azure( container="container", connection_string="conn", blob_name="name", secrets=["foo"], labels=["bar", "baz"], add_default_labels=False, )
def test_azure_creds_not_serialized(): azure = storage.Azure(container="container", connection_string="conn", blob_name="name") serialized = AzureSchema().dump(azure) assert serialized assert serialized["__version__"] == prefect.__version__ assert serialized["container"] == "container" assert serialized["blob_name"] == "name" assert serialized.get("connection_string") is None
def test_azure_full_serialize(): azure = storage.Azure( container="container", connection_string="conn", blob_name="name", secrets=["foo"], labels=["bar", "baz"], add_default_labels=False, ) serialized = AzureSchema().dump(azure) assert serialized assert serialized["__version__"] == prefect.__version__ assert serialized["container"] == "container" assert serialized["blob_name"] == "name" assert serialized["secrets"] == ["foo"]
def configure_flow_storage(cluster: Cluster, secrets): if cluster.flow_storage_protocol == S3_PROTOCOL: key = secrets[cluster.flow_storage_options.key] secret = secrets[cluster.flow_storage_options.secret] flow_storage = storage.S3( bucket=cluster.flow_storage, client_options={ "aws_access_key_id": key, "aws_secret_access_key": secret }, ) return flow_storage elif cluster.flow_storage_protocol == ABFS_PROTOCOL: secret = secrets[cluster.flow_storage_options.secret] flow_storage = storage.Azure(container=cluster.flow_storage, connection_string=secret) return flow_storage else: raise UnsupportedFlowStorage
def test_azure_serialize_with_flows(): azure = storage.Azure( container="container", connection_string="conn", blob_name="name", secrets=["foo"], ) f = prefect.Flow("test") azure.flows["test"] = "key" serialized = AzureSchema().dump(azure) assert serialized assert serialized["__version__"] == prefect.__version__ assert serialized["container"] == "container" assert serialized["blob_name"] == "name" assert serialized.get("connection_string") is None assert serialized["flows"] == {"test": "key"} deserialized = AzureSchema().load(serialized) assert f.name in deserialized assert deserialized.secrets == ["foo"]
def register_recipe(recipe: BaseRecipe): fs_remote = AzureBlobFileSystem( connection_string=os.environ["FLOW_STORAGE_CONNECTION_STRING"] ) target = FSSpecTarget( fs_remote, root_path=f"abfs://{os.environ['FLOW_CACHE_CONTAINER']}/azurerecipetest/", ) recipe.target = target recipe.input_cache = CacheFSSpecTarget( fs_remote, root_path=( f"abfs://{os.environ['FLOW_CACHE_CONTAINER']}/azurerecipetestcache/" ), ) recipe.metadata_cache = target executor = PrefectPipelineExecutor() pipeline = recipe.to_pipelines() flow = executor.pipelines_to_plan(pipeline) job_template = yaml.safe_load( """ apiVersion: batch/v1 kind: Job metadata: annotations: "cluster-autoscaler.kubernetes.io/safe-to-evict": "false" spec: template: spec: containers: - name: flow """ ) flow_name = "test-noaa-flow" flow.storage = storage.Azure( container=os.environ["FLOW_STORAGE_CONTAINER"], connection_string=os.environ["FLOW_STORAGE_CONNECTION_STRING"], ) flow.run_config = KubernetesRun( job_template=job_template, image=os.environ["BAKERY_IMAGE"], env={ "AZURE_STORAGE_CONNECTION_STRING": os.environ[ "FLOW_STORAGE_CONNECTION_STRING" ], }, labels=json.loads(os.environ["PREFECT__CLOUD__AGENT__LABELS"]), cpu_request="1000m", memory_request="3Gi", ) worker_spec = make_pod_spec( image=os.environ["BAKERY_IMAGE"], labels={"flow": flow_name}, memory_limit="1Gi", memory_request="500Mi", cpu_limit="512m", cpu_request="256m", env={ "AZURE_STORAGE_CONNECTION_STRING": os.environ[ "FLOW_STORAGE_CONNECTION_STRING" ] }, ) scheduler_spec = make_pod_spec( image=os.environ["BAKERY_IMAGE"], labels={"flow": flow_name}, memory_request="500Mi", cpu_request="256m", ) scheduler_spec.spec.containers[0].args = ["dask-scheduler"] scheduler_spec = clean_pod_template(scheduler_spec, pod_type="scheduler") flow.executor = DaskExecutor( cluster_class="dask_kubernetes.KubeCluster", cluster_kwargs={ "pod_template": worker_spec, "scheduler_pod_template": scheduler_spec, }, adapt_kwargs={"maximum": 10}, ) for flow_task in flow.tasks: flow_task.run = set_log_level(flow_task.run) flow.name = flow_name project_name = os.environ["PREFECT_PROJECT"] flow_id = flow.register(project_name=project_name) return flow_id
def test_azure_creds_not_serialized(): azure = storage.Azure( container="container", connection_string="conn", blob_name="name" )
def test_azure_empty_serialize(): azure = storage.Azure(container="container")