def test_import_missing_dependency(): with pytest.raises( ImportError, match=".*%s.*PURPOSE.*pip install 'bionic\\[%s\\]'.*" % (TEST_PACKAGE_NAME, TEST_EXTRA_NAME), ): import_optional_dependency(TEST_PACKAGE_NAME, purpose="PURPOSE")
def loky_executor(): loky = import_optional_dependency("loky", purpose="parallel execution") return loky.get_reusable_executor( max_workers=None, initializer=logging_initializer, initargs=(get_singleton_manager().logging_queue,), )
def _stage(self): cloudpickle = import_optional_dependency("cloudpickle") path = self.inputs_uri() logging.info(f"Staging task {self.name} at {path}") with self.gcs_fs.open(path, "wb") as f: cloudpickle.dump(self, f)
def get_aip_client(cache_value=True): if cache_value: global _cached_aip_client if _cached_aip_client is None: _cached_aip_client = get_aip_client(cache_value=False) return _cached_aip_client discovery = import_optional_dependency("googleapiclient.discovery", raise_on_missing=True) logger.info("Initializing AIP client ...") return discovery.build("ml", "v1", cache_discovery=False)
def _run(ipath, gcs_fs): cloudpickle = import_optional_dependency("cloudpickle") with gcs_fs.open(ipath, "rb") as f: task = cloudpickle.load(f) # Now that we have the task, set up logging. _set_up_logging(task.job_id(), task.config.project_name) logging.info(f"Read task from {ipath}") result = task.function() opath = task.output_uri() logging.info(f"Uploading result to {opath}") with gcs_fs.open(opath, "wb") as f: pickle.dump(result, f)
def _set_up_logging(job_id, project_id): if os.environ.get("BIONIC_NO_STACKDRIVER", False): return # TODO This is the ID of the hyperparameter tuning trial currently # running on this VM. This field is only set if the current # training job is a hyperparameter tuning job. Conductor uses this # environment variable but AIP documentation suggests us to use # TF_CONFIG. Check whether we need to update this env variable. # Find more details on TF_CONFIG at this link: # https://cloud.google.com/ai-platform/training/docs/distributed-training-details trial_id = os.environ.get("CLOUD_ML_TRIAL_ID", None) glogging = import_optional_dependency("google.cloud.logging") client = glogging.Client(project=project_id) resource = glogging.resource.Resource( type="ml_job", # AIP expects a default task_name for the master cluster. We # use a placeholder value till we start using clusters. Once we # do, it should be configured based on the cluster. labels=dict(job_id=job_id, project_id=project_id, task_name="master-replica-0"), ) labels = None if trial_id is not None: # Enable grouping by trial when present. labels = {"ml.googleapis.com/trial_id": trial_id} # Enable only the cloud logger to avoid duplicate messages. handler = glogging.handlers.handlers.CloudLoggingHandler(client, resource=resource, labels=labels) root_logger = logging.getLogger() # Remote the StreamHandler. Any logs logged by it shows up as error # logs in Stackdriver. root_logger.handlers = [] # We should ideally make this configurable, but till then, let's # set the level to DEBUG to write all the logs. It's not hard to # filter using log level on Stackdriver so it doesn't create too # much noise anyway. root_logger.setLevel(logging.DEBUG) root_logger.addHandler(handler) for logger_name in glogging.handlers.handlers.EXCLUDED_LOGGER_DEFAULTS: logging.getLogger(logger_name).propagate = False
def test_import_unrecognized_dependency(): with pytest.raises(AssertionError): import_optional_dependency("_UNKNOWN_PACKAGE_", purpose="PURPOSE")
def test_import_missing_dependency_without_raising(): module = import_optional_dependency(TEST_PACKAGE_NAME, raise_on_missing=False) assert module is None
def get_gcp_project_id(): google_auth = import_optional_dependency( "google.auth", purpose="Get GCP project id from the environment") _, project = google_auth.default() return project